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/Cargo.lock b/Cargo.lock index 98323720b..7d2a04292 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -5689,6 +5689,7 @@ dependencies = [ "humantime", "indicatif", "jsonwebtoken", + "nix", "object_store", "parquet", "pgwire-replication", @@ -5727,6 +5728,7 @@ dependencies = [ "unicode-properties", "uuid", "webpki-roots", + "windows-sys 0.61.2", "zip", "zstd", ] diff --git a/Cargo.toml b/Cargo.toml index ba18d3060..3cf5d4451 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -49,6 +49,7 @@ hmac = "0.12" humantime = "2" console = { version = "0.16", default-features = false } indicatif = { version = "0.18", default-features = false, features = ["unicode-width"] } +nix = { version = "0.31", default-features = false, features = ["resource"] } object_store = { version = "0.14", default-features = false, features = ["aws-base", "reqwest", "ring", "fs"] } jsonwebtoken = { version = "11", default-features = false, features = ["use_pem", "rust_crypto"] } parquet = { version = "60", default-features = false, features = ["arrow", "zstd", "async", "object_store", "snap", "lz4", "flate2", "flate2-rust_backend"] } @@ -92,11 +93,12 @@ tracing-subscriber = { version = "0.3", features = ["fmt", "json", "env-filter"] unicode-properties = { version = "0.1.5", default-features = false, features = ["general-category"] } uuid = { version = "1", features = ["v4"] } webpki-roots = "1" +windows-sys = { version = "0.61", features = ["Win32_Foundation", "Win32_System_ProcessStatus", "Win32_System_Threading"] } zstd = "0.14" zip = { version = "9", default-features = false, features = ["deflate"] } [workspace.lints.rust] -unsafe_code = "forbid" +unsafe_code = "deny" unreachable_pub = "deny" unused_must_use = "deny" 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/Cargo.toml b/crates/sqlload/Cargo.toml index 441bc19e4..d60a02c08 100644 --- a/crates/sqlload/Cargo.toml +++ b/crates/sqlload/Cargo.toml @@ -116,12 +116,14 @@ webpki-roots.workspace = true zstd.workspace = true [target.'cfg(any(target_os = "linux", target_os = "macos"))'.dependencies] +nix.workspace = true rustix.workspace = true # librdkafka's default build runs its mklove `configure` script and `make`, which MSVC has # neither of; on Windows it is built with CMake instead. [target.'cfg(windows)'.dependencies] rdkafka = { workspace = true, features = ["cmake-build"] } +windows-sys.workspace = true [dev-dependencies] bytes.workspace = true 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/files/main/partition_window.rs b/crates/sqlload/src/adapters/files/main/partition_window.rs index 35899573e..bd133e680 100644 --- a/crates/sqlload/src/adapters/files/main/partition_window.rs +++ b/crates/sqlload/src/adapters/files/main/partition_window.rs @@ -16,7 +16,7 @@ pub fn partition_window(request: &PartitionWindowRequest<'_>) -> Result ( Some(align_down(start, by, cursor_type)?), - format!("{column} from {start} (--start-cursor)"), + format!("{column} from {start} ({})", request.start_flag), ), (None, Some(newest), _) => { let start = partition_start(newest, by, cursor_type); diff --git a/crates/sqlload/src/adapters/files/models.rs b/crates/sqlload/src/adapters/files/models.rs index 88afd8b6e..b120f4b85 100644 --- a/crates/sqlload/src/adapters/files/models.rs +++ b/crates/sqlload/src/adapters/files/models.rs @@ -150,6 +150,8 @@ pub struct PartitionWindowRequest<'a> { pub lookback: Option, /// `--start-cursor-*`, already normalised for the cursor type. pub start: Option<&'a str>, + /// The flag `start` came from, as `plan` names it: `--start-cursor`, `--start-cursor-ts` or `--start-cursor-int`. + pub start_flag: &'a str, /// `--end-cursor-*`, already normalised for the cursor type. pub end: Option<&'a str>, pub cursor_start: Option<&'a str>, diff --git a/crates/sqlload/src/adapters/files/tests/test_partition_window.rs b/crates/sqlload/src/adapters/files/tests/test_partition_window.rs index babec1e51..6e5af3988 100644 --- a/crates/sqlload/src/adapters/files/tests/test_partition_window.rs +++ b/crates/sqlload/src/adapters/files/tests/test_partition_window.rs @@ -140,6 +140,7 @@ fn given_existing_partitions_when_choosing_the_window_then_it_covers_whole_parti newest: test_case.newest, lookback: test_case.lookback, start: test_case.flags.0, + start_flag: "--start-cursor", end: test_case.flags.1, cursor_start: test_case.cursor_start, reload: test_case.reload, 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/cdc/_helpers/stream.rs b/crates/sqlload/src/cdc/_helpers/stream.rs index a50ece563..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(); @@ -1285,6 +1303,7 @@ pub(crate) fn source_outcomes( retryable: loaded.is_err() && retryable, elapsed, skipped: false, + peak_memory_bytes: None, }); } outcomes @@ -1461,6 +1480,7 @@ pub(crate) fn failed_outcomes( retryable: false, elapsed, skipped: false, + peak_memory_bytes: None, }], false => outcomes, } 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..570ccdc2a 100644 --- a/crates/sqlload/src/command/_helpers/load_output.rs +++ b/crates/sqlload/src/command/_helpers/load_output.rs @@ -2,17 +2,23 @@ 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; 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::memory_text::memory_text; 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; @@ -47,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)] @@ -94,6 +100,8 @@ pub(crate) struct Tally { pub(crate) rows: u64, /// The table that took longest to load, and how long. pub(crate) slowest: Option<(String, u64)>, + /// The most memory the whole process held at once, in bytes, as the run's end reported it. + pub(crate) peak_memory_bytes: Option, } impl Tally { @@ -116,6 +124,9 @@ impl Tally { } EventKind::TableFailed { .. } => self.failed += 1, EventKind::TableNotStarted { .. } => self.not_started.push(table), + EventKind::RunFinished { + peak_memory_bytes, .. + } => self.peak_memory_bytes = *peak_memory_bytes, _ => {} } } @@ -130,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. @@ -147,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 { @@ -160,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(); @@ -172,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, .. } => { @@ -231,16 +266,28 @@ pub(crate) fn event_lines( elapsed_ms, position, changes, + transforms, + peak_memory_bytes, } => { let word = match rows { 1 => "row", _ => "rows", }; let seconds = format!("{:.1}s", *elapsed_ms as f64 / 1000.0); + let note: Vec = position + .iter() + .cloned() + .chain(peak_memory_bytes.map(peak_note)) + .collect(); 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), + note.join(" ") + ) + .trim_end(), + transform_lines(transforms) ); let mut lines: Vec = changes .iter() @@ -249,12 +296,66 @@ pub(crate) fn event_lines( lines.push(tally_line(("loaded", name), &detail, name_width)); return lines; } - EventKind::TableFailed { reason, .. } => ("failed", reason.clone()), + 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, + .. + } => ( + "failed", + match peak_memory_bytes { + Some(bytes) => format!("{reason}\n{}", peak_note(*bytes)), + None => reason.clone(), + }, + ), _ => return Vec::new(), }; vec![tally_line((verb, name), &detail, name_width)] } +/// A table's peak memory under `--memory-per-table`, as its tally line notes it: `peak memory 85 MiB`. +fn peak_note(bytes: u64) -> String { + format!("peak memory {}", memory_text(bytes)) +} + +/// 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); @@ -307,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)); @@ -343,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( @@ -360,7 +477,7 @@ fn live_bar() -> indicatif::ProgressBar { bar } -/// The line a load ends with, and the branches under it: what never started, the slowest table, and after failures whether running it again helps and what it would do. +/// The line a load ends with, and the branches under it: what never started, the slowest table, the run's peak memory, and after failures whether running it again helps and what it would do. pub(crate) fn summary_lines(tally: &Tally, elapsed_ms: u64, exit_code: u8) -> Vec { let mut clauses: Vec = Vec::new(); clauses.extend((tally.loaded > 0).then(|| format!("{} loaded", tally.loaded))); @@ -397,6 +514,9 @@ pub(crate) fn summary_lines(tally: &Tally, elapsed_ms: u64, exit_code: u8) -> Ve format!("{table} {:.1}s", *elapsed_ms as f64 / 1000.0), )); } + if let Some(bytes) = tally.peak_memory_bytes { + branches.push(("peak memory", memory_text(bytes))); + } if failures && exit_code == crate::command::constants::RETRYABLE_EXIT { branches.push(( "retry", @@ -503,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, @@ -514,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, @@ -599,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 fdf3f0318..69ccb043a 100644 --- a/crates/sqlload/src/command/_helpers/loader_commands.rs +++ b/crates/sqlload/src/command/_helpers/loader_commands.rs @@ -3,8 +3,10 @@ use std::io::IsTerminal; use std::path::Path; use std::sync::Arc; +use crate::adapters::contract::main::cursor_text::cursor_text; use crate::adapters::dispatch::models::Destination; use crate::auth::main::allow_sign_in::allow_sign_in; +use crate::column_types::models::ColumnType; use crate::command::_helpers::load_output::{ EventsTarget, LiveState, LoadOutput, change_tally_line, planning, start_lines, summary_lines, write_events, @@ -12,7 +14,9 @@ use crate::command::_helpers::load_output::{ use crate::command::_helpers::output::{confirm_reload, print_lines}; use crate::command::_helpers::project_file::replace_file; use crate::command::_helpers::signals::wait_for_signal; -use crate::command::constants::{FORCED_STOP_EXIT, RELOAD_DECLINED, RETRYABLE_EXIT}; +use crate::command::constants::{ + CURSOR_FLAGS_MIXED, FORCED_STOP_EXIT, RELOAD_DECLINED, RETRYABLE_EXIT, +}; use crate::command::errors::CommandError; use crate::command::models::{ Command, ContractAction, EventsFormat, LoaderArgs, LoaderRequest, Progress, StateAction, @@ -25,7 +29,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; @@ -36,9 +40,12 @@ use crate::loader::main::render_plan::render_plan; use crate::loader::main::render_tally::render_tally; use crate::loader::main::state_command::state_command; use crate::loader::models::{ - LoadRun, PlanOptions, ProjectPlan, RunReport, StateKind, StateRequest, TableOutcome, TablePlan, + LoadRun, PlanOptions, ProjectPlan, RangeFlags, RunReport, StateKind, StateRequest, + TableOutcome, TablePlan, TypedCursors, }; use crate::number_text::main::thousands::thousands; +use crate::peak_memory::constants::MEMORY_PER_TABLE_NEEDS_LINUX; +use crate::peak_memory::main::run_peak_bytes::run_peak_bytes; use crate::project::main::config_lines::config_lines; use crate::project::main::parse_project::parse_project; use crate::project::main::parse_python_project::parse_python_project; @@ -63,16 +70,61 @@ fn variable(variables: &BTreeMap, name: &str) -> Option /// The plan options `args` give. pub(crate) fn options(args: LoaderArgs) -> Result { + if args.memory_per_table && !cfg!(target_os = "linux") { + return Err(MEMORY_PER_TABLE_NEEDS_LINUX.to_owned()); + } + let typed_cursors = TypedCursors { + start_ts: args.start_cursor_ts, + end_ts: args.end_cursor_ts, + start_int: args.start_cursor_int, + end_int: args.end_cursor_int, + }; + let typed = [ + ( + "--start-cursor-ts", + &typed_cursors.start_ts, + ColumnType::TimestampTz, + ), + ( + "--end-cursor-ts", + &typed_cursors.end_ts, + ColumnType::TimestampTz, + ), + ( + "--start-cursor-int", + &typed_cursors.start_int, + ColumnType::Int, + ), + ("--end-cursor-int", &typed_cursors.end_int, ColumnType::Int), + ]; + let untyped = args.start_cursor.is_some() || args.end_cursor.is_some(); + if untyped && typed.iter().any(|(_, value, _)| value.is_some()) { + return Err(CURSOR_FLAGS_MIXED.to_owned()); + } + for (flag, value, column_type) in &typed { + let Some(value) = value else { + continue; + }; + if cursor_text(value, column_type).is_err() { + return Err(match column_type { + ColumnType::Int => format!("{flag} value '{value}' is not a whole number"), + _ => format!("{flag} value '{value}' is not an RFC 3339 timestamp or a date"), + }); + } + } Ok(PlanOptions { + typed_cursors, select: [args.positional_select, args.select].concat(), exclude: args.exclude, rebuild_allowed: args.allow_type_change_rebuild, reload: args.reload, start_cursor: args.start_cursor, end_cursor: args.end_cursor, + range_flags: RangeFlags::default(), concurrency: args.concurrency.map(usize::from), resnapshot: args.resnapshot, fail_fast: args.fail_fast, + memory_per_table: args.memory_per_table, batch_bytes: args .batch_bytes .map(|bytes| usize::try_from(bytes).unwrap_or(usize::MAX)), @@ -98,6 +150,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 +166,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. @@ -310,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)) } @@ -354,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), @@ -372,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, @@ -382,58 +475,60 @@ async fn run_command( if streamed.output.print { print_lines(std::mem::take(&mut lines), no_color); } - let (outcomes, kept, code) = Box::pin(streamed_load(streamed)).await; + let ended = Box::pin(streamed_load(streamed)).await; if json { let run = RunReport { command: "load", applied: &applied, - outcomes: &outcomes, + outcomes: &ended.outcomes, started: started.0, elapsed: started.1.elapsed(), + peak_memory_bytes: ended.peak_memory_bytes, }; let (text, ok, retryable) = json_report(&plan, Some(&run)); return Ok((vec![text], exit_code(ok, retryable))); } - lines.extend(kept); - return Ok((lines, code)); + 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(&[]), }; 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()); + 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), } -/// Load `load.plan`, rendering its events as they come: each table's outcome, the lines not printed yet (the `created`/`altered` lines `apply` reported, each table's line as it finished, and the summary, which a load of no table leaves to the plan's closing line), and the exit code, the same as a load without events. -async fn streamed_load(load: StreamedLoad<'_>) -> (Vec, Vec, u8) { +/// How a streamed load ended: each table's outcome, the lines not printed yet, the exit code, and the run's peak memory in bytes. +struct StreamedEnd { + outcomes: Vec, + lines: Vec, + exit_code: u8, + peak_memory_bytes: Option, +} + +/// Load `load.plan`, rendering its events as they come: each table's outcome, the lines not printed yet (the `created`/`altered` lines `apply` reported, each table's line as it finished, and the summary, which a load of no table leaves to the plan's closing line), the exit code, the same as a load without events, and the run's peak memory (the most of the process's before the tables loaded, after, and each table's under `--memory-per-table`, whose resets the process's own figure may miss). +async fn streamed_load(load: StreamedLoad<'_>) -> StreamedEnd { let receiver = listen(&uuid::Uuid::new_v4().to_string()); let width = load.output.name_width; let selected = load.plan.tables.len(); @@ -461,7 +556,12 @@ async fn streamed_load(load: StreamedLoad<'_>) -> (Vec, Vec> = std::iter::once(before) + .chain(outcomes.iter().map(|outcome| outcome.peak_memory_bytes)) + .collect(); + let peak_memory_bytes = run_peak_bytes(&measured); let (mut loaded, mut failed, mut not_started, mut rows) = (0, 0, 0, 0); for outcome in &outcomes { match (&outcome.loaded, outcome.skipped) { @@ -484,6 +584,7 @@ async fn streamed_load(load: StreamedLoad<'_>) -> (Vec, Vec) -> (Vec, Vec> = 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. 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_writer(std::io::stderr); + .with_ansi(coloured) + .with_writer(|| LogWriter); let _ = match format { LogFormat::Text => builder.try_init(), LogFormat::Json => builder.json().try_init(), diff --git a/crates/sqlload/src/command/constants.rs b/crates/sqlload/src/command/constants.rs index 51050f08f..ca925dd99 100644 --- a/crates/sqlload/src/command/constants.rs +++ b/crates/sqlload/src/command/constants.rs @@ -23,3 +23,7 @@ pub const WRAP_ROOM: usize = 20; /// The narrowest terminal a line is wrapped for; a narrower one gets each line as it is. pub const MIN_WRAP_WIDTH: usize = 2; + +/// Why `--start-cursor`/`--end-cursor` and the typed cursor flags are refused together, in SQLBuild's words. +pub const CURSOR_FLAGS_MIXED: &str = + "--start-cursor/--end-cursor cannot be combined with the -ts or -int cursor flags"; diff --git a/crates/sqlload/src/command/models.rs b/crates/sqlload/src/command/models.rs index 55b54b11b..5c8dd7bda 100644 --- a/crates/sqlload/src/command/models.rs +++ b/crates/sqlload/src/command/models.rs @@ -210,6 +210,18 @@ pub struct LoaderArgs { /// Backfill only cursor values below this one, typed from each selected table's cursor. #[arg(long, value_name = "VALUE")] pub end_cursor: Option, + /// Backfill the selected date and timestamp cursor tables from this RFC 3339 timestamp or date; not with --start-cursor/--end-cursor. + #[arg(long, value_name = "TIMESTAMP")] + pub start_cursor_ts: Option, + /// Backfill the selected date and timestamp cursor tables only below this RFC 3339 timestamp or date. + #[arg(long, value_name = "TIMESTAMP")] + pub end_cursor_ts: Option, + /// Backfill the selected integer cursor tables from this integer; not with --start-cursor/--end-cursor. + #[arg(long, value_name = "INTEGER", allow_hyphen_values = true)] + pub start_cursor_int: Option, + /// Backfill the selected integer cursor tables only below this integer. + #[arg(long, value_name = "INTEGER", allow_hyphen_values = true)] + pub end_cursor_int: Option, /// `load`: once a table fails, start no other table; tables already loading finish, and the run exits 1, or 75 when the failure is retryable. #[arg(long, conflicts_with = "follow")] pub fail_fast: bool, @@ -220,7 +232,7 @@ pub struct LoaderArgs { #[arg(long)] pub json: bool, /// `load` of one postgres_cdc source only: keep streaming after catching up, writing a batch every `batch_interval`, until --max-duration passes or Ctrl-C/SIGTERM (a second one stops at once). - #[arg(long, conflicts_with_all = ["reload", "start_cursor", "end_cursor"])] + #[arg(long, conflicts_with_all = ["reload", "start_cursor", "end_cursor", "start_cursor_ts", "end_cursor_ts", "start_cursor_int", "end_cursor_int"])] pub follow: bool, /// With --follow: stop after this long, e.g. `10m`, once the buffered changes are written. #[arg(long, requires = "follow", value_parser = humantime::parse_duration)] @@ -234,6 +246,9 @@ pub struct LoaderArgs { /// `load`: draw the live block of running tables on stderr: on a terminal without --json or --log-format json (`auto`), always, or never. #[arg(long, value_enum, default_value_t = Progress::Auto)] pub progress: Progress, + /// `load`, Linux only: load one table at a time and report each table's peak memory (the run's peak is always reported). + #[arg(long, conflicts_with = "follow")] + pub memory_per_table: bool, /// The bytes a batch holds before it may commit, instead of about 1 GB (for tests). #[arg(long, hide = true, value_parser = clap::value_parser!(u64).range(1..))] pub batch_bytes: Option, diff --git a/crates/sqlload/src/command/tests/helpers.rs b/crates/sqlload/src/command/tests/helpers.rs index 6e2e2fd95..3c1e99b47 100644 --- a/crates/sqlload/src/command/tests/helpers.rs +++ b/crates/sqlload/src/command/tests/helpers.rs @@ -91,6 +91,8 @@ 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(), + peak_memory_bytes: None, } } diff --git a/crates/sqlload/src/command/tests/test_load_options.rs b/crates/sqlload/src/command/tests/test_load_options.rs index 823ecb4ac..8bc9758da 100644 --- a/crates/sqlload/src/command/tests/test_load_options.rs +++ b/crates/sqlload/src/command/tests/test_load_options.rs @@ -20,9 +20,33 @@ fn given_selection_and_cursor_flags_when_parsing_a_load_then_the_plan_options_ho expected_options: Ok(" | | 100 | 2026-11-05"), }, LoadOptionsTestCase { - description: "the typed cursor flags are gone", - argv: &["--start-cursor-int", "7"], - expected_options: Err("unexpected argument '--start-cursor-int'"), + description: "the -ts and -int flags combine, as in SQLBuild, each for its own cursor family", + argv: &[ + "--start-cursor-ts", + "2026-10-01", + "--end-cursor-int", + "2027", + ], + expected_options: Ok(" | | | "), + }, + LoadOptionsTestCase { + description: "the untyped and typed flags are refused together, in SQLBuild's words", + argv: &["--start-cursor", "100", "--end-cursor-int", "200"], + expected_options: Err( + "--start-cursor/--end-cursor cannot be combined with the -ts or -int cursor flags", + ), + }, + LoadOptionsTestCase { + description: "an -int value must be a whole number", + argv: &["--start-cursor-int", "3.14"], + expected_options: Err("--start-cursor-int value '3.14' is not a whole number"), + }, + LoadOptionsTestCase { + description: "a -ts value must be an RFC 3339 timestamp or a date", + argv: &["--end-cursor-ts", "yesterday"], + expected_options: Err( + "--end-cursor-ts value 'yesterday' is not an RFC 3339 timestamp or a date", + ), }, LoadOptionsTestCase { description: "--threads is gone; --concurrency replaces it", diff --git a/crates/sqlload/src/command/tests/test_load_output.rs b/crates/sqlload/src/command/tests/test_load_output.rs index 43daeac0a..53b32b2da 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,8 @@ 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(), + peak_memory_bytes: None, }, )], name_width: 7, @@ -64,6 +69,85 @@ 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", + events: vec![ + ( + Some("raw.a"), + EventKind::TableDone { + rows: 1, + elapsed_ms: 300, + position: Some("id now 7".to_owned()), + changes: Vec::new(), + transforms: Vec::new(), + peak_memory_bytes: Some(85 * 1024 * 1024), + }, + ), + ( + Some("raw.b"), + EventKind::TableFailed { + reason: "HTTP 404 not found".to_owned(), + retryable: false, + elapsed_ms: 100, + peak_memory_bytes: Some(1_932_735_283), + }, + ), + ( + None, + EventKind::RunFinished { + loaded: 1, + failed: 1, + not_started: 0, + rows: 1, + elapsed_ms: 500, + exit_code: 1, + peak_memory_bytes: Some(2_000_000_000), + }, + ), + ], + name_width: 7, + rows_width: 5, + elapsed_ms: 500, + exit_code: 1, + expected_lines: "loaded raw.a 1 row 0.3s id now 7 peak memory 85 MiB +failed raw.b HTTP 404 not found + peak memory 1.8 GiB +✗ Completed with failures 1 loaded, 1 failed, 1 row, 0.5s +├── slowest raw.a 0.3s +├── peak memory 1.9 GiB +└── next the 1 unchanged table loads; the 1 that committed carries on from its position", + }, ]; for test_case in test_cases { let mut tally = Tally::default(); @@ -132,3 +216,75 @@ 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(), + peak_memory_bytes: None, + }, + }; + 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..9eb32ab37 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, @@ -74,6 +79,14 @@ fn given_a_plan_that_fails_when_planning_then_it_prints_the_errors_or_one_json_e expected_output: "Plan ready 0 selected", expected_exit: "exit 0", }, + #[cfg(not(target_os = "linux"))] + PlanFailureTestCase { + description: "--memory-per-table resets Linux's high-water mark, so plan refuses it elsewhere", + project: TARGET, + argv: &["plan", "--memory-per-table"], + expected_output: "--memory-per-table needs Linux; the run's peak is in every report", + expected_exit: "exit 1", + }, ]; for test_case in test_cases { let output = run_without_host(test_case.project, test_case.argv); 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/contract/tests/helpers.rs b/crates/sqlload/src/contract/tests/helpers.rs index a827ad00d..e6af41230 100644 --- a/crates/sqlload/src/contract/tests/helpers.rs +++ b/crates/sqlload/src/contract/tests/helpers.rs @@ -82,7 +82,7 @@ pub(crate) fn python_spec(settings: &str, warehouse: Warehouse) -> Result 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/job.rs b/crates/sqlload/src/http_source/_helpers/job.rs index add7cd135..0ebecb9e4 100644 --- a/crates/sqlload/src/http_source/_helpers/job.rs +++ b/crates/sqlload/src/http_source/_helpers/job.rs @@ -125,7 +125,7 @@ fn context_problems(job: &HttpJob) -> Vec { } /// Why `key_path`'s `path` can't be read, if it can't: it is a field path below `$`, such as `$.issueId` or `$.keys[0]`. -fn key_path_problem(text: &str) -> Option { +pub(crate) fn key_path_problem(text: &str) -> Option { match field_path(text) { Ok(segments) if !segments.is_empty() => None, Ok(_) => Some(format!( 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/main/literal_path_problems.rs b/crates/sqlload/src/http_source/main/literal_path_problems.rs new file mode 100644 index 000000000..63bf0293d --- /dev/null +++ b/crates/sqlload/src/http_source/main/literal_path_problems.rs @@ -0,0 +1,30 @@ +//! Check the `sl.http` paths a Python table writes as literals with load's own checker, before any request. + +use crate::http_source::_helpers::job::key_path_problem; +use crate::http_source::models::LiteralPaths; +use crate::row_batches::main::projection::projection; +use crate::row_batches::models::FieldSpec; + +/// Why load would refuse each of `paths`, in load's words: `fields["sku"]: field path "$.lines[*].sku" uses the wildcard [*], ...`, `key_path ...`, or `transform kpis: no function is registered as kpis ...`. +pub fn literal_path_problems(paths: &LiteralPaths) -> Vec { + let mut problems: Vec = Vec::new(); + for (name, path) in &paths.fields { + let field = FieldSpec { + name: name.clone(), + path: path.clone(), + required: false, + parse: None, + on_error: Default::default(), + from: Default::default(), + }; + problems.extend(projection(Some(&[field]), None).err()); + } + problems.extend( + paths + .key_paths + .iter() + .filter_map(|path| key_path_problem(path)), + ); + problems.extend(paths.transform_problems.iter().cloned()); + problems +} diff --git a/crates/sqlload/src/http_source/main/mod.rs b/crates/sqlload/src/http_source/main/mod.rs index 264153a58..f48db0ec0 100644 --- a/crates/sqlload/src/http_source/main/mod.rs +++ b/crates/sqlload/src/http_source/main/mod.rs @@ -1 +1,2 @@ +pub mod literal_path_problems; pub mod run_http_job; diff --git a/crates/sqlload/src/http_source/models.rs b/crates/sqlload/src/http_source/models.rs index 2231c2858..9ea40f2be 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. @@ -455,3 +479,36 @@ pub enum Sink { Pages(Sender), Keys(Sender), } + +/// The `sl.http` paths a Python table function writes as literals, which discovery reads from its source without calling it, for `plan` to check as load would. +#[derive(Debug, Clone, PartialEq, Eq, Default, Deserialize)] +#[serde(default, deny_unknown_fields)] +pub struct LiteralPaths { + /// Each `fields=` entry of a job that loads rows, as its column name and path. + pub fields: Vec<(String, String)>, + /// Each `key_path=` of a job that loads rows. + pub key_paths: Vec, + /// What `plan` can't read as a literal, such as `fields["sku"]` or `key_path`: load checks it. + pub computed: Vec, + /// The table function's source can't be read, so load checks every path. + pub no_source: bool, + /// Why each `transform=` written as a name can't be looked up: no `@sl.transform` function, or two, registered under it. + pub transform_problems: Vec, +} + +impl LiteralPaths { + /// The fact `plan` shows when load checks some of the paths: those that are computed, or all of them when the source can't be read. + pub fn checked_at_load(&self) -> Option { + if self.no_source { + return Some( + "checked at load, as the table function's source can't be read".to_owned(), + ); + } + (!self.computed.is_empty()).then(|| { + format!( + "checked at load, as they're computed: {}", + self.computed.join(", ") + ) + }) + } +} diff --git a/crates/sqlload/src/lib.rs b/crates/sqlload/src/lib.rs index 89fc88b2f..a5219ed0c 100644 --- a/crates/sqlload/src/lib.rs +++ b/crates/sqlload/src/lib.rs @@ -1,6 +1,6 @@ -//! SQLLoad: the Kafka to Snowflake exporter and the loader engine. +//! SQLLoad: the Kafka to Snowflake exporter and the loader engine. It denies `unsafe` rather than forbidding it so that one Windows-only helper may allow it: reading the process's peak memory needs an FFI call no maintained crate wraps safely. -#![forbid(unsafe_code)] +#![deny(unsafe_code)] pub mod adapters; pub mod auth; @@ -36,6 +36,7 @@ pub mod number_text; #[cfg(feature = "odbc")] pub mod odbc_source; pub mod parquet_parts; +pub mod peak_memory; pub mod planning; pub mod postgres_source; pub mod project; diff --git a/crates/sqlload/src/load_events/models.rs b/crates/sqlload/src/load_events/models.rs index 80df8bb51..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, @@ -43,16 +49,25 @@ 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, + /// The most memory the process held while the table loaded, in bytes; only under `--memory-per-table`. + #[serde(default, skip_serializing_if = "Option::is_none")] + peak_memory_bytes: Option, }, /// A table failed, and whether running the load again may get past it. TableFailed { reason: String, retryable: bool, elapsed_ms: u64, + /// The most memory the process held while the table loaded, in bytes; only under `--memory-per-table`. + #[serde(default, skip_serializing_if = "Option::is_none")] + peak_memory_bytes: Option, }, /// A table the load never started, and why. TableNotStarted { reason: String }, - /// The load ended: its counts and the exit code. + /// The load ended: its counts, the exit code, and the most memory the whole process held at once. RunFinished { loaded: usize, failed: usize, @@ -60,5 +75,8 @@ pub enum EventKind { rows: u64, elapsed_ms: u64, exit_code: u8, + /// In bytes; absent where the OS gives no figure. + #[serde(default, skip_serializing_if = "Option::is_none")] + peak_memory_bytes: Option, }, } 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..ea7c806fb 100644 --- a/crates/sqlload/src/loader/_helpers/python.rs +++ b/crates/sqlload/src/loader/_helpers/python.rs @@ -22,6 +22,8 @@ use crate::column_types::main::arrow_field::arrow_field; use crate::column_types::main::combined_type::combined_type; use crate::column_types::main::exact_cast::exact_cast; use crate::column_types::models::ColumnType; +use crate::http_source::main::literal_path_problems::literal_path_problems; +use crate::http_source::models::LiteralPaths; use crate::load_naming::constants::CHILD_TABLE_SEPARATOR; use crate::load_naming::main::child_table_name::child_table_name; use crate::load_naming::main::loaded_name::loaded_name; @@ -39,6 +41,7 @@ use crate::loader::_helpers::table_plan::replan_columns; use crate::loader::_helpers::table_plan::schema_statements; use crate::loader::constants::DROPPED_CHUNK; use crate::loader::constants::EMPTY_REPLACE_PROBLEM; +use crate::loader::constants::ENFORCED_UNDECLARED_PROBLEM; use crate::loader::constants::NO_DATA_NOTE; use crate::loader::constants::NO_STATE; use crate::loader::constants::PYTHON_COLUMN_TYPE; @@ -67,6 +70,7 @@ use crate::postgres_source::models::SourceColumn; use crate::postgres_source::models::SourceTable; use crate::project::constants::FILES_CONNECTION_KINDS; use crate::project::models::CheckName; +use crate::project::models::Contract; use crate::project::models::NullCursor; use crate::project::models::OnSchemaChange; use crate::project::models::ProjectConfig; @@ -151,8 +155,12 @@ pub(crate) fn python_tables(source: &SourceConfig) -> Vec { tables } -/// What a Python table's spec records about it. -pub(crate) fn python_table(module: &str, table: &SourceTable) -> PythonTable { +/// What a Python table's spec records about it, with the `sl.http` paths its function writes as literals among `paths`, by table name. +pub(crate) fn python_table( + module: &str, + table: &SourceTable, + paths: Option<&BTreeMap>, +) -> PythonTable { PythonTable { module: module.to_owned(), declared: table @@ -161,9 +169,23 @@ pub(crate) fn python_table(module: &str, table: &SourceTable) -> PythonTable { .map(|column| column.name.clone()) .collect(), lists: ListChoices::new(), + literal_paths: paths.and_then(|paths| paths.get(&table.name)).cloned(), } } +/// Why a Python table can't load, found before it is planned: an enforced contract with no declared columns, or an `sl.http` path its function writes as a literal that load would refuse. +pub(crate) fn python_problems(spec: &TableSpec) -> Vec { + let Some(python) = &spec.python else { + return Vec::new(); + }; + let mut problems: Vec = Vec::new(); + if python.declared.is_empty() && spec.contract == Contract::Enforced { + problems.push(ENFORCED_UNDECLARED_PROBLEM.to_owned()); + } + problems.extend(python.literal_paths.iter().flat_map(literal_path_problems)); + problems +} + /// The child tables `spec`'s table has in the destination (tables named `__...` 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, @@ -686,7 +708,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 +721,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 +927,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 +951,7 @@ impl PythonFeed { note: None, after_load: None, replaced_keys, + transforms: Vec::new(), }) } @@ -1353,8 +1380,10 @@ 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, cdc: None, checks: None, warehouse_source: None, @@ -1388,6 +1417,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 cb56394a1..df128afec 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)); @@ -619,6 +673,16 @@ 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())); + } if matches!(table.commits, CommitRule::Windows(_)) && spec.python.is_none() && spec.files.is_none() @@ -648,6 +712,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 +1094,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 +1104,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 +1147,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 +1159,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 +1224,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(), @@ -1258,6 +1347,7 @@ pub(crate) fn run_result(plan: &ProjectPlan, run: &RunReport<'_>) -> RunResult { outcomes, started, elapsed, + peak_memory_bytes, } = *run; let changes = outcomes .iter() @@ -1294,6 +1384,7 @@ pub(crate) fn run_result(plan: &ProjectPlan, run: &RunReport<'_>) -> RunResult { failed: statuses("failed"), skipped: statuses("skipped"), rows: tables.iter().filter_map(|table| table.rows).sum(), + peak_memory_bytes, }; let view = plan_view(plan); let finished = started + chrono::Duration::from_std(elapsed).unwrap_or_default(); @@ -1344,6 +1435,11 @@ 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(), + }, + peak_memory_bytes: outcome.peak_memory_bytes, } } diff --git a/crates/sqlload/src/loader/_helpers/resolve.rs b/crates/sqlload/src/loader/_helpers/resolve.rs index c4ac85251..c153cd331 100644 --- a/crates/sqlload/src/loader/_helpers/resolve.rs +++ b/crates/sqlload/src/loader/_helpers/resolve.rs @@ -1,24 +1,37 @@ 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::column_types::models::Warehouse; +use crate::adapters::dispatch::main::session::live_columns::live_columns; +use crate::adapters::dispatch::models::Destination; +use crate::column_types::models::{ColumnType, 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::{ + DATE_TEXT_LENGTH, NO_TRANSFORM_HOST, PRIMARY_KEY_PLACEHOLDER, RANGE_NEEDS_CURSOR, +}; +use crate::loader::errors::LoaderError; +use crate::loader::models::{ + PlanOptions, ProjectPlan, RangeFlags, ResolvedTables, SkippedTable, TablePlan, TableSpec, + TableTransform, TypedCursors, 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, NullCursor, ProjectConfig, + SourceConfig, SourceKind, TableOverride, TableSettings, UnnestSetting, }; +use crate::schema_plan::models::SchemaPlan; +use crate::similar_names::main::nearest_name::nearest_name; use crate::ssh_tunnel::models::SshTunnel; use std::collections::BTreeMap; @@ -64,7 +77,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 +150,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 +597,361 @@ 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() +} + +/// The cursors a typed range flag applies to: dates and timestamps (`-ts`), or integers (`-int`). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum CursorFamily { + Time, + Integer, +} + +/// A typed range flag given: its name, its value and the cursors it applies to. +type GivenFlag<'a> = (&'static str, &'a String, CursorFamily); + +/// One table's `options` and `schema` once its cursor, `cursor_type` (`None` while a Python table's cursor is unseen), is known: its family's typed range as `start_cursor`/`end_cursor`, and `schema` blocked by why the typed flags or the lookback can't apply, or warned of the rows a ranged `--reload` drops. +pub(crate) fn ranged_table( + options: &PlanOptions, + spec: &TableSpec, + cursor_type: Option<&ColumnType>, + mut schema: SchemaPlan, +) -> (PlanOptions, SchemaPlan) { + let (ranged, problem) = match table_range(options, spec, cursor_type, false) { + Ok(ranged) => (ranged, None), + Err(problem) => (options.clone(), Some(problem)), + }; + if let Some(problem) = problem.filter(|_| schema.blocked.is_none()) { + schema = blocked_plan(vec![problem]); + } + schema.warnings.extend(reload_range(spec, &ranged)); + (ranged, schema) +} + +/// `options` for one table whose cursor is `cursor_type` (`None` while a Python table's cursor is unseen), into a files target when `files_target`: with typed flags, its family's range as `start_cursor`/`end_cursor`; or why the typed flags or its `lookback` can't apply to it. +pub(crate) fn table_range( + options: &PlanOptions, + spec: &TableSpec, + cursor_type: Option<&ColumnType>, + files_target: bool, +) -> Result { + let ranged = table_options(options, spec, cursor_type)?; + lookback_problem(spec, cursor_type, files_target).map_or(Ok(ranged), Err) +} + +/// The typed flags given, in the order `plan` names them. +fn given_flags(typed: &TypedCursors) -> Vec> { + let flags = [ + ("--start-cursor-ts", &typed.start_ts, CursorFamily::Time), + ("--end-cursor-ts", &typed.end_ts, CursorFamily::Time), + ( + "--start-cursor-int", + &typed.start_int, + CursorFamily::Integer, + ), + ("--end-cursor-int", &typed.end_int, CursorFamily::Integer), + ]; + let mut given: Vec> = Vec::new(); + for (flag, value, family) in flags { + if let Some(value) = value { + given.push((flag, value, family)); + } + } + given +} + +/// [`table_range`] without the lookback. +fn table_options( + options: &PlanOptions, + spec: &TableSpec, + cursor_type: Option<&ColumnType>, +) -> Result { + let typed = &options.typed_cursors; + let given = given_flags(typed); + if given.is_empty() { + return Ok(options.clone()); + } + let Some(cursor) = &spec.cursor else { + return Err(RANGE_NEEDS_CURSOR.to_owned()); + }; + let named: Vec<&str> = given.iter().map(|(flag, _, _)| *flag).collect(); + let has = |family: CursorFamily| given.iter().any(|(_, _, given)| *given == family); + let family = match cursor_type { + Some(ColumnType::Int | ColumnType::SmallInt) => CursorFamily::Integer, + Some(ColumnType::Date | ColumnType::TimestampTz | ColumnType::TimestampNtz) => { + CursorFamily::Time + } + Some(other) => { + return Err(format!( + "cursor {} is {other}, which takes --start-cursor/--end-cursor, not {}: the -ts flags are for date and timestamp cursors and the -int flags for integer ones", + cursor.label(), + named.join(", ") + )); + } + None => match (has(CursorFamily::Time), has(CursorFamily::Integer)) { + (true, false) => CursorFamily::Time, + (false, true) => CursorFamily::Integer, + _ => { + return Err(format!( + "cursor {} has no type until the table first loads, so it can't choose between the -ts and -int flags; give one family's flags, or --start-cursor/--end-cursor", + cursor.label() + )); + } + }, + }; + if !has(family) { + let (given_kind, wanted) = match family { + CursorFamily::Time => ("integer", "--start-cursor-ts/--end-cursor-ts"), + CursorFamily::Integer => ("date and timestamp", "--start-cursor-int/--end-cursor-int"), + }; + return Err(format!( + "{} {} for {given_kind} cursors, but cursor {} is {}; give this table {wanted}, or leave it out of the selection", + named.join(", "), + match named.len() { + 1 => "is", + _ => "are", + }, + cursor.label(), + cursor_type.map(ToString::to_string).unwrap_or_default(), + )); + } + if cursor_type == Some(&ColumnType::Date) + && let Some((flag, value, _)) = given.iter().find(|(_, value, family)| { + *family == CursorFamily::Time && value.trim().len() > DATE_TEXT_LENGTH + }) + { + return Err(format!( + "{flag} {value} has a time of day, but cursor {} is a date; give a date, such as {}", + cursor.label(), + value.trim().get(..DATE_TEXT_LENGTH).unwrap_or_default() + )); + } + let (start, end, range_flags) = match family { + CursorFamily::Time => ( + &typed.start_ts, + &typed.end_ts, + RangeFlags { + start: "--start-cursor-ts", + end: "--end-cursor-ts", + }, + ), + CursorFamily::Integer => ( + &typed.start_int, + &typed.end_int, + RangeFlags { + start: "--start-cursor-int", + end: "--end-cursor-int", + }, + ), + }; + Ok(PlanOptions { + start_cursor: start.clone(), + end_cursor: end.clone(), + range_flags, + ..options.clone() + }) +} + +/// Why `spec`'s `lookback` can't apply to its cursor, `cursor_type` (`None` while a Python table's cursor is unseen), or into a files target (`files_target`): an integer lookback that counts back no units, or on anything but an integer cursor; a duration on anything but a date or timestamp cursor. +fn lookback_problem( + spec: &TableSpec, + cursor_type: Option<&ColumnType>, + files_target: bool, +) -> Option { + let cursor = spec.cursor.as_ref()?; + let label = cursor.label(); + if let Some(units) = cursor.lookback_units { + if units < 1 { + return Some(format!( + "lookback = {units} counts back no cursor units; give 1 or more, or leave lookback out" + )); + } + if files_target { + return Some(format!( + "lookback = {units} counts back integer cursor units, which a files target's window can't: it follows its newest partition; leave lookback out" + )); + } + return match cursor_type? { + ColumnType::Int | ColumnType::SmallInt => None, + time @ (ColumnType::Date | ColumnType::TimestampTz | ColumnType::TimestampNtz) => { + Some(format!( + "lookback = {units} counts integer cursor units, but cursor {label} is {time}; use a duration, such as lookback = \"1d\"" + )) + } + other => Some(format!( + "lookback needs a date, timestamp or integer cursor, but cursor {label} is {other}; leave lookback out" + )), + }; + } + let duration = cursor.lookback?; + match cursor_type? { + ColumnType::Date | ColumnType::TimestampTz | ColumnType::TimestampNtz => None, + ColumnType::Int | ColumnType::SmallInt => Some(format!( + "lookback = \"{}\" is a duration, but cursor {label} is an integer; use a number of cursor units, such as lookback = 1", + humantime::format_duration(duration) + )), + other => Some(format!( + "lookback needs a date, timestamp or integer cursor, but cursor {label} is {other}; leave lookback out" + )), + } +} + +/// What a `--reload` with a cursor range (`options`, as one table's [`table_range`] gave them) leaves out of the rebuilt table: the rows outside that range are dropped; rows with a NULL cursor are kept under `null_cursor = "include"`, which reads them with the first window, and dropped under `"ignore"` (under `"fail"`, the load fails on them). +fn reload_range(spec: &TableSpec, options: &PlanOptions) -> Option { + let RangeFlags { + start: start_flag, + end: end_flag, + } = options.range_flags; + let (start, end) = ( + options.start_cursor.as_deref().filter(|_| options.reload), + options.end_cursor.as_deref().filter(|_| options.reload), + ); + let (kept, dropped) = match (start, end) { + (Some(start), Some(end)) => ( + format!("from {start} ({start_flag}), below {end} ({end_flag})"), + format!("rows before {start} and from {end} on are dropped"), + ), + (Some(start), None) => ( + format!("from {start} ({start_flag})"), + format!("rows before {start} are dropped"), + ), + (None, Some(end)) => ( + format!("below {end} ({end_flag})"), + format!("rows from {end} on are dropped"), + ), + (None, None) => return None, + }; + let nulls = match spec.cursor.as_ref().map(|cursor| cursor.null_cursor) { + Some(NullCursor::Include) => "; rows with a NULL cursor are kept", + Some(NullCursor::Ignore) => "; rows with a NULL cursor are dropped", + Some(NullCursor::Fail) | None => "", + }; + Some(format!( + "--reload rebuilds the table with only the rows {kept}: {dropped}{nulls}" + )) +} diff --git a/crates/sqlload/src/loader/_helpers/run_log.rs b/crates/sqlload/src/loader/_helpers/run_log.rs index 561897864..ced518132 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; @@ -283,7 +291,7 @@ pub(crate) fn run_entry( let logged = load.table.logged.as_ref()?; let spec = &table.spec; let options = load.options; - let backfill = options.start_cursor.is_some() || options.end_cursor.is_some(); + let backfill = options.ranged(); let previous = logged .current .as_ref() @@ -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..f36c73a1a 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 { .. })) @@ -359,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()), @@ -368,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( @@ -565,6 +632,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 +706,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 +784,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 +815,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 +845,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 +905,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 +918,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 +1141,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 +1168,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 +1178,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 +1231,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 +1278,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 +1299,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 +1316,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 +1373,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 +1436,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 }; @@ -1450,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?) } @@ -1468,14 +1494,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 +1511,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, @@ -1524,7 +1649,10 @@ pub(crate) async fn source_ends( mut specs: Vec, options: &PlanOptions, ) -> Result, LoaderError> { - if source.kind != SourceKind::Postgres || options.end_cursor.is_some() { + if source.kind != SourceKind::Postgres + || options.end_cursor.is_some() + || options.typed_cursors.end_int.is_some() + { return Ok(specs); } let mut client: Option = None; @@ -1551,31 +1679,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 +1696,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 +1877,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 +1949,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..500594d85 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, + CdcDeletes, CdcMode, CursorBounds, Lookback, LookbackFrom, MaskMethod, NullCursor, NullKeys, + 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)?); @@ -799,7 +818,14 @@ pub(crate) fn cursor(table: &TableSettings, source: &TableSettings) -> Option Some(duration), + _ => None, + }, + lookback_units: match layered(&table.lookback, &source.lookback) { + Some((Lookback::Units(units), _)) => Some(units), + _ => None, + }, lookahead: layered(&table.lookahead, &source.lookahead).map(|(value, _)| value), from_now: layered(&table.lookback_from, &source.lookback_from) .is_some_and(|(value, _)| value == LookbackFrom::Now), @@ -1093,10 +1119,14 @@ pub(crate) fn strategy_problems( (WriteStrategy::Merge, _) if unique_key.is_empty() => problems.push(format!( "write_strategy = \"merge\" needs a unique_key: the key columns, such as unique_key = [\"id\"], or unique_key = {PRIMARY_KEY_PLACEHOLDER:?} for the primary key the source reports" )), - (WriteStrategy::Append, Some(cursor)) if cursor.lookback.is_some() => problems.push( - "lookback re-inserts the window on every run with append; use delete_insert (no key needed) or merge" - .to_owned(), - ), + (WriteStrategy::Append, Some(cursor)) + if cursor.lookback.is_some() || cursor.lookback_units.is_some() => + { + problems.push( + "lookback re-inserts the window on every run with append; use delete_insert (no key needed) or merge" + .to_owned(), + ) + } (WriteStrategy::Replace, Some(_)) => problems.push( "write_strategy = \"replace\" reloads the whole table, so it takes no cursor_column".to_owned(), ), @@ -1114,7 +1144,16 @@ pub(crate) fn strategy_problems( "replace_key is for write_strategy = \"delete_insert\", which replaces the rows holding the key values a load reads; {strategy} does not delete by it. Fix: set write_strategy = \"delete_insert\" to replace those rows whole, or drop replace_key{merge}" )); } - if let Some(cursor) = cursor.filter(|cursor| cursor.from_now && cursor.lookback.is_none()) { + if let Some(cursor) = cursor.filter(|cursor| cursor.from_now && cursor.lookback_units.is_some()) + { + problems.push(format!( + "lookback_from = \"now\" counts lookback back from the time of the run, so {} needs a duration lookback, such as lookback = \"1d\", not a number of cursor units", + cursor.label() + )); + } + if let Some(cursor) = cursor.filter(|cursor| { + cursor.from_now && cursor.lookback.is_none() && cursor.lookback_units.is_none() + }) { problems.push(format!( "lookback_from = \"now\" counts lookback back from now, so {} needs a lookback", cursor.label() @@ -1630,3 +1669,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..663498c3e 100644 --- a/crates/sqlload/src/loader/_helpers/table_plan.rs +++ b/crates/sqlload/src/loader/_helpers/table_plan.rs @@ -41,20 +41,23 @@ 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::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; use crate::loader::_helpers::resolve::display_name; +use crate::loader::_helpers::resolve::ranged_table; +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; +use crate::loader::constants::RANGE_NEEDS_CURSOR; use crate::loader::errors::LoaderError; use crate::loader::main::_files_schema::files_schema; use crate::loader::main::_logged_plan::with_run_log; @@ -75,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}; @@ -219,6 +221,7 @@ pub(crate) fn blocked_table(spec: &TableSpec, reasons: Vec) -> TablePlan commits: CommitRule::Whole(WholeReason::Blocked), stage: None, planned_stage: None, + fetch_memory: None, } } @@ -228,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?, @@ -251,9 +248,9 @@ pub(crate) async fn plan_table( None => plan, }; 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.reload || options.ranged()).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. @@ -377,37 +374,6 @@ fn split_ranking(plan: &TablePlan) -> Option { )) } -/// What a `--reload` with `--start-cursor` or `--end-cursor` leaves out of the rebuilt table: the rows outside that range are dropped; rows with a NULL cursor are kept under `null_cursor = "include"`, which reads them with the first window, and dropped under `"ignore"` (under `"fail"`, the load fails on them). -fn reload_range(spec: &TableSpec, options: &PlanOptions) -> Option { - let (start, end) = ( - options.start_cursor.as_deref().filter(|_| options.reload), - options.end_cursor.as_deref().filter(|_| options.reload), - ); - let (kept, dropped) = match (start, end) { - (Some(start), Some(end)) => ( - format!("from {start} (--start-cursor), below {end} (--end-cursor)"), - format!("rows before {start} and from {end} on are dropped"), - ), - (Some(start), None) => ( - format!("from {start} (--start-cursor)"), - format!("rows before {start} are dropped"), - ), - (None, Some(end)) => ( - format!("below {end} (--end-cursor)"), - format!("rows from {end} on are dropped"), - ), - (None, None) => return None, - }; - let nulls = match spec.cursor.as_ref().map(|cursor| cursor.null_cursor) { - Some(NullCursor::Include) => "; rows with a NULL cursor are kept", - Some(NullCursor::Ignore) => "; rows with a NULL cursor are dropped", - Some(NullCursor::Fail) | None => "", - }; - Some(format!( - "--reload rebuilds the table with only the rows {kept}: {dropped}{nulls}" - )) -} - /// Why the table's `batch_size` setting has no effect, if it has none: a files source commits by whole partitions, `replace` swaps the whole table in once, as does `--reload` (whose cursor windows still cut the reads), and windows need a cursor. fn ignored_batch(plan: &TablePlan) -> Option { let spec = &plan.spec; @@ -787,6 +753,7 @@ async fn plan_table_once( commits: CommitRule::Whole(WholeReason::Blocked), stage: None, planned_stage: None, + fetch_memory: None, }); } }; @@ -822,7 +789,13 @@ async fn plan_table_once( false => cursor_type, }; let cursor_name = spec.cursor.as_ref().map(|cursor| cursor.column.as_str()); - schema = checked_cursor(spec, schema, (&desired, unseen_python_cursor)); + let (ranged, mut schema) = ranged_table( + options, + spec, + cursor_type.filter(|_| !unseen_python_cursor), + checked_cursor(spec, schema, (&desired, unseen_python_cursor)), + ); + let options = &ranged; let cursor_in_table = has_cursor(spec, live.as_deref()); let cursor_expression = source_cursor(spec, &read_columns, cursor_type); let lacks_cursor = @@ -979,6 +952,7 @@ async fn plan_table_once( commits, stage, planned_stage, + fetch_memory: None, }) } @@ -1019,8 +993,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 +1065,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 +1239,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, @@ -1432,6 +1424,24 @@ fn is_number(value: &str) -> bool { && digits(exponent) } +/// Where an integer cursor's read starts under `lookback = units`: the destination's `maximum` minus `units`, but never below the smallest value `cursor_type` holds (`smallint` is 16-bit, `int` 64-bit), which the read's bound must fit. +pub(crate) fn lookback_start( + maximum: &str, + units: i64, + cursor_type: &ColumnType, +) -> Result { + let lowest = match cursor_type { + ColumnType::SmallInt => i64::from(i16::MIN), + _ => i64::MIN, + }; + match maximum.trim().parse::() { + Ok(value) => Ok(value.saturating_sub(units).max(lowest).to_string()), + Err(_) => Err(format!( + "cursor maximum {maximum:?} is not an integer, so lookback = {units} can't count back from it" + )), + } +} + /// Where the read starts, and why: a `--start-cursor`, the destination's maximum (read after it by `append` from a database or Python source, unless `append_cursor_inclusive`; a files source reads from it, skipping files already loaded), `cursor_start`, or nothing; with the destination's maximum (without lookback) when it was read. async fn lower_bound( request: &WindowRequest<'_>, @@ -1443,7 +1453,11 @@ async fn lower_bound( if let Some(start) = start_flag { return Ok(Ok(( Some(start.to_owned()), - format!("{} from {start} (--start-cursor)", cursor.label()), + format!( + "{} from {start} ({})", + cursor.label(), + options.range_flags.start + ), None, ))); } @@ -1523,6 +1537,13 @@ async fn lower_bound( Some(Err(problem)) => return Ok(Err(problem)), None => (None, None), }; + let (maximum, lookback) = match (maximum, cursor.lookback_units) { + (Some(maximum), Some(units)) => match lookback_start(&maximum, units, cursor_type) { + Ok(start) => (Some(start), format!(" {maximum} minus lookback {units}")), + Err(problem) => return Ok(Err(problem)), + }, + (maximum, _) => (maximum, lookback), + }; Ok(Ok(match (maximum, &cursor_start) { (Some(maximum), _) => { let bound = match request.spec.strategy == WriteStrategy::Append @@ -1557,12 +1578,9 @@ pub(crate) async fn cursor_window( ) -> Result, LoaderError> { let spec = request.spec; let options = request.options; - let backfill = options.start_cursor.is_some() || options.end_cursor.is_some(); let (Some(cursor), Some(cursor_type)) = (&spec.cursor, request.cursor_type) else { - if backfill { - return Ok(Err( - "--start-cursor/--end-cursor need a cursor_column".to_owned() - )); + if options.ranged() { + return Ok(Err(RANGE_NEEDS_CURSOR.to_owned())); } return Ok(Ok(Window { filter: None, @@ -1609,7 +1627,7 @@ pub(crate) async fn cursor_window( let (to, note) = match (to, ahead) { (Some(end), _) => ( Some(end.clone()), - format!("{note}, below {end} (--end-cursor)"), + format!("{note}, below {end} ({})", options.range_flags.end), ), (None, Some((end, lookahead))) => ( Some(end.clone()), @@ -1866,6 +1884,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/_helpers/table_queue.rs b/crates/sqlload/src/loader/_helpers/table_queue.rs index 9f73895fc..1692c66ac 100644 --- a/crates/sqlload/src/loader/_helpers/table_queue.rs +++ b/crates/sqlload/src/loader/_helpers/table_queue.rs @@ -2,18 +2,24 @@ 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; use crate::loader::constants::DEFAULT_THREADS; use crate::loader::main::_load_table::load_one; use crate::loader::models::{LoadRun, ProjectPlan, SharedRun, TableOutcome, TablePlan}; +use crate::peak_memory::main::with_table_peak::with_table_peak; use crate::project::models::ProjectConfig; /// The most tables loading at once. pub(crate) fn thread_count(run: &LoadRun<'_>) -> usize { + if run.options.memory_per_table { + return 1; + } run.options .concurrency .or(run.project.concurrency) @@ -40,7 +46,12 @@ pub(crate) async fn load_task(shared: Arc, table: TablePlan) -> Table options: &shared.options, }; let started = std::time::Instant::now(); - match load_one(&run, &table).await { + let (loaded, peak_memory_bytes) = with_table_peak( + shared.options.memory_per_table, + Box::pin(load_one(&run, &table)), + ) + .await; + match loaded { Ok((name, loaded)) => TableOutcome { table: name, origin: Some(table.origin()), @@ -48,6 +59,7 @@ pub(crate) async fn load_task(shared: Arc, table: TablePlan) -> Table retryable: false, elapsed: started.elapsed(), skipped: false, + peak_memory_bytes, }, Err(error) => TableOutcome { table: table.outcome_name(), @@ -56,10 +68,34 @@ pub(crate) async fn load_task(shared: Arc, table: TablePlan) -> Table loaded: Err(error.to_string()), elapsed: started.elapsed(), skipped: false, + peak_memory_bytes, }, } } +/// 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 { @@ -81,6 +117,8 @@ pub(crate) fn emit_outcome(outcome: &TableOutcome) { elapsed_ms, position: loaded.note.clone(), changes: loaded.changes.clone(), + transforms: loaded.transforms.clone(), + peak_memory_bytes: outcome.peak_memory_bytes, }, Err(reason) if outcome.skipped => EventKind::TableNotStarted { reason: reason.clone(), @@ -89,6 +127,7 @@ pub(crate) fn emit_outcome(outcome: &TableOutcome) { reason: reason.clone(), retryable: outcome.retryable, elapsed_ms, + peak_memory_bytes: outcome.peak_memory_bytes, }, }; emit(Some(&outcome.table), kind); @@ -217,6 +256,7 @@ impl<'a> TableQueue<'a> { retryable: false, elapsed: std::time::Duration::ZERO, skipped: false, + peak_memory_bytes: None, }); emit_outcome(&outcome); if self.fail_fast && self.stopped.is_none() && outcome.loaded.is_err() { diff --git a/crates/sqlload/src/loader/constants.rs b/crates/sqlload/src/loader/constants.rs index fd42a769e..cb91368ba 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,13 @@ 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 length of a date written as `YYYY-MM-DD`; a longer `--start-cursor-ts` value has a time of day. +pub const DATE_TEXT_LENGTH: usize = 10; + +/// Why a cursor range can't apply to a table without a cursor. +pub const RANGE_NEEDS_CURSOR: &str = + "--start-cursor/--end-cursor and their -ts and -int forms need a cursor_column"; + +/// 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/_files_window.rs b/crates/sqlload/src/loader/main/_files_window.rs index 62c6cc9fc..967463a81 100644 --- a/crates/sqlload/src/loader/main/_files_window.rs +++ b/crates/sqlload/src/loader/main/_files_window.rs @@ -41,7 +41,7 @@ pub(in crate::loader) fn files_window(request: &FilesWindowRequest<'_>) -> Resul }; if backfill && options.reload { return Err( - "--reload rewrites every file, so it takes no --start-cursor/--end-cursor".to_owned(), + "--reload rewrites every file, so it takes no --start-cursor/--end-cursor, nor their -ts and -int forms".to_owned(), ); } if cursor.lookahead.is_some() || cursor.from_now { @@ -76,6 +76,7 @@ pub(in crate::loader) fn files_window(request: &FilesWindowRequest<'_>) -> Resul newest, lookback: cursor.lookback, start: start.as_deref(), + start_flag: options.range_flags.start, end: end.as_deref(), cursor_start: cursor.start.as_deref(), reload: options.reload, @@ -99,7 +100,7 @@ pub(in crate::loader) fn files_window(request: &FilesWindowRequest<'_>) -> Resul let (from, note, operator) = match (start, maximum, &cursor.start) { (Some(start), _, _) => ( Some(start.clone()), - format!("{column} from {start} (--start-cursor)"), + format!("{column} from {start} ({})", options.range_flags.start), INCLUSIVE_BOUND, ), (None, Some((maximum, note)), _) => { diff --git a/crates/sqlload/src/loader/main/_logged_plan.rs b/crates/sqlload/src/loader/main/_logged_plan.rs index 7627e35c1..c936af6c5 100644 --- a/crates/sqlload/src/loader/main/_logged_plan.rs +++ b/crates/sqlload/src/loader/main/_logged_plan.rs @@ -7,15 +7,15 @@ use crate::loader::_helpers::run_log::{ }; use crate::loader::_helpers::table_plan::blocked_plan; use crate::loader::errors::LoaderError; -use crate::loader::models::{LoggedPlan, PlanOptions, TablePlan}; +use crate::loader::models::{LoggedPlan, TablePlan}; use crate::run_log::main::current_run::current_run; use crate::run_log::models::RunEvent; -/// `plan` with its run log when the table is logged: after `sqlload logged-state reset`, the next run reads from the reset position (unless flags choose the window); otherwise a logged cursor that disagrees with the data's maximum is noted, and the data wins. A record that cannot be read blocks the table until a reset follows it. +/// `plan` with its run log when the table is logged: after `sqlload logged-state reset`, the next run reads from the reset position (unless `flags`, `--reload` or a cursor range, choose the window); otherwise a logged cursor that disagrees with the data's maximum is noted, and the data wins. A record that cannot be read blocks the table until a reset follows it. pub(in crate::loader) async fn with_run_log( mut plan: TablePlan, destination: &Destination, - options: &PlanOptions, + flags: bool, ) -> Result { if !plan.spec.logged || plan.schema.blocked.is_some() { return Ok(plan); @@ -32,7 +32,6 @@ pub(in crate::loader) async fn with_run_log( }); let records = log.records; let current = current_run(&records).cloned(); - let flags = options.reload || options.start_cursor.is_some() || options.end_cursor.is_some(); let reset = current .as_ref() .filter(|record| record.event == RunEvent::Reset && !flags); diff --git a/crates/sqlload/src/loader/main/_plan_files.rs b/crates/sqlload/src/loader/main/_plan_files.rs index 1c9b5cc3e..6bd24ef8e 100644 --- a/crates/sqlload/src/loader/main/_plan_files.rs +++ b/crates/sqlload/src/loader/main/_plan_files.rs @@ -7,6 +7,7 @@ use crate::adapters::files::models::FilesDestination; use crate::file_store::main::open_folder::open_folder; use crate::files_source::models::{FilesRead, LoadedFile}; use crate::loader::_helpers::destination::{live_schema, table_folder}; +use crate::loader::_helpers::resolve::table_range; use crate::loader::_helpers::sources::{FilesPlanRequest, LoadedSource, plan_files}; use crate::loader::_helpers::table_plan::{CommitRequest, Window, commit_rule}; use crate::loader::_helpers::table_plan::{ @@ -106,6 +107,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(), @@ -118,7 +120,7 @@ pub(in crate::loader) async fn plan_files_table( }) } -/// The window `request` reads in (from the target's newest partition or files), and for a files source the files it reads, skipping those the target's records hold (`loaded`), with no source filter, as the files are chosen instead; a problem in either blocks `schema` instead. Nothing is planned for a plan already blocked. +/// The window `request` reads in (from the target's newest partition or files, over the range of the cursor's family when typed range flags are given), and for a files source the files it reads, skipping those the target's records hold (`loaded`), with no source filter, as the files are chosen instead; a problem in either blocks `schema` instead. Nothing is planned for a plan already blocked. async fn planned_read( request: &FilesWindowRequest<'_>, (loaded, live): (&[LoadedFile], Option<&[LiveColumn]>), @@ -127,6 +129,14 @@ async fn planned_read( if schema.blocked.is_some() { return Ok((schema, None, None)); } + let ranged = match table_range(request.options, request.spec, request.cursor_type, true) { + Ok(ranged) => ranged, + Err(problem) => return Ok((blocked_plan(vec![problem]), None, None)), + }; + let request = &FilesWindowRequest { + options: &ranged, + ..*request + }; let window = match files_window(request) { Ok(window) => window, Err(problem) => return Ok((blocked_plan(vec![problem]), None, None)), 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/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/crates/sqlload/src/loader/main/plan_project.rs b/crates/sqlload/src/loader/main/plan_project.rs index 183f86fac..4a6a6ce39 100644 --- a/crates/sqlload/src/loader/main/plan_project.rs +++ b/crates/sqlload/src/loader/main/plan_project.rs @@ -10,11 +10,13 @@ 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, }; #[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; use crate::loader::_helpers::table_plan::{blocked_plan, plan_table, skip_up_to_date}; use crate::loader::errors::LoaderError; @@ -77,7 +79,7 @@ pub async fn plan_project( .await?; PLANNED_SOURCES.fetch_add(1, Ordering::Relaxed); } - let ranged = options.start_cursor.is_some() || options.end_cursor.is_some(); + let ranged = options.ranged(); Ok(skip_up_to_date( block_conflicts(plan, (&claims, &configured, &drained), (project, options)), ranged, @@ -92,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; @@ -110,6 +113,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 +130,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/main/render_tally.rs b/crates/sqlload/src/loader/main/render_tally.rs index c007be103..49920148e 100644 --- a/crates/sqlload/src/loader/main/render_tally.rs +++ b/crates/sqlload/src/loader/main/render_tally.rs @@ -4,13 +4,15 @@ use std::time::Duration; use crate::loader::_helpers::render::change_line; use crate::loader::models::{TableOutcome, TallyLine}; +use crate::number_text::main::memory_text::memory_text; use crate::number_text::main::thousands::thousands; -/// One aligned line per object `applied` created or altered and per table loaded, failed or skipped, then a final line counting them; and whether nothing failed. +/// One aligned line per object `applied` created or altered and per table loaded, failed or skipped, then a final line counting them, and under it the run's peak memory when a load gives `peak_memory_bytes`; and whether nothing failed. pub fn render_tally( applied: &[String], outcomes: &[TableOutcome], elapsed: Duration, + peak_memory_bytes: Option, ) -> (Vec, bool) { let mut lines: Vec = applied.iter().map(|line| change_line(line)).collect(); let mut rows: u64 = 0; @@ -75,5 +77,6 @@ pub fn render_tally( 0 => format!("✓ Completed {}", clauses.join(", ")), _ => format!("✗ Completed with failures {}", clauses.join(", ")), }); + out.extend(peak_memory_bytes.map(|bytes| format!("└── peak memory {}", memory_text(bytes)))); (out, failed == 0) } diff --git a/crates/sqlload/src/loader/main/selected_specs.rs b/crates/sqlload/src/loader/main/selected_specs.rs index fde879dcf..fb391c626 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}; @@ -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 9166a4bc6..e27e7d78e 100644 --- a/crates/sqlload/src/loader/models.rs +++ b/crates/sqlload/src/loader/models.rs @@ -59,7 +59,10 @@ impl std::fmt::Display for SettingOrigin { pub struct CursorSpec { pub column: String, pub start: Option, + /// A duration `lookback`, on a date or timestamp cursor. pub lookback: Option, + /// An integer `lookback`, in cursor units, on an integer cursor. + pub lookback_units: Option, /// Read only cursor values before now plus this (`lookahead`). pub lookahead: Option, /// `lookback_from = "now"`: incremental runs read from now minus `lookback`, not from the destination's maximum minus it. @@ -201,6 +204,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. @@ -263,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`. @@ -292,6 +312,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. @@ -404,6 +441,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. @@ -604,10 +658,14 @@ pub struct PlanOptions { pub rebuild_allowed: bool, /// `--reload`: rebuild each selected table from the whole source, ignoring the destination's cursor. pub reload: bool, - /// `--start-cursor` (or `-ts`/`-int`): read from this cursor value instead of the destination's maximum. + /// `--start-cursor` (or, for one table, its family's `-ts`/`-int`): read from this cursor value instead of the destination's maximum. pub start_cursor: Option, - /// `--end-cursor` (or `-ts`/`-int`): read only cursor values below this one. + /// `--end-cursor` (or, for one table, its family's `-ts`/`-int`): read only cursor values below this one. pub end_cursor: Option, + /// `--start-cursor-ts`/`--end-cursor-ts` and `--start-cursor-int`/`--end-cursor-int`: each cursor table reads the range of its own family. + pub typed_cursors: TypedCursors, + /// The flags `start_cursor` and `end_cursor` came from, as `plan` names them: the untyped ones, or the typed ones of the family one table's options took. + pub range_flags: RangeFlags, /// `--concurrency`: the most tables loading at once, overriding the project's `concurrency`. pub concurrency: Option, /// The bytes a batch holds before it may commit, instead of about 1 GB (a hidden flag, for tests). @@ -616,6 +674,50 @@ pub struct PlanOptions { pub resnapshot: bool, /// `--fail-fast`: once a table fails, start no other; tables already loading finish. pub fail_fast: bool, + /// `--memory-per-table` (Linux): one table at a time, each table's peak memory measured. + pub memory_per_table: bool, +} + +/// The typed cursor range flags, as given. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct TypedCursors { + pub start_ts: Option, + pub end_ts: Option, + pub start_int: Option, + pub end_int: Option, +} + +impl PlanOptions { + /// Whether a cursor range is given, by `--start-cursor`/`--end-cursor` or a typed flag. + pub fn ranged(&self) -> bool { + let typed = &self.typed_cursors; + [ + &self.start_cursor, + &self.end_cursor, + &typed.start_ts, + &typed.end_ts, + &typed.start_int, + &typed.end_int, + ] + .iter() + .any(|value| value.is_some()) + } +} + +/// The flags a cursor range came from, as `plan` names them. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct RangeFlags { + pub start: &'static str, + pub end: &'static str, +} + +impl Default for RangeFlags { + fn default() -> Self { + Self { + start: "--start-cursor", + end: "--end-cursor", + } + } } /// What one Python table's load needs. @@ -659,6 +761,8 @@ pub struct TableOutcome { pub elapsed: std::time::Duration, /// Never started, because `--fail-fast` stopped the load after another table failed; `loaded` holds why. pub skipped: bool, + /// The most memory the process held while the table loaded, in bytes; only under `--memory-per-table`. + pub peak_memory_bytes: Option, } impl TableOutcome { @@ -673,6 +777,7 @@ impl TableOutcome { retryable: true, elapsed: std::time::Duration::ZERO, skipped: true, + peak_memory_bytes: None, } } @@ -704,6 +809,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 +950,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 +1017,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. @@ -1089,6 +1224,8 @@ pub struct RunReport<'a> { pub outcomes: &'a [TableOutcome], pub started: chrono::DateTime, pub elapsed: std::time::Duration, + /// The run's peak memory, in bytes, when the OS gives it. + pub peak_memory_bytes: Option, } /// What `apply --json` and `load --json` print, in this field order. @@ -1122,6 +1259,9 @@ pub struct RunSummary { /// Tables `--fail-fast` never started. pub skipped: usize, pub rows: u64, + /// The most memory the whole process held at once during the run, in bytes; left out where the OS gives no figure. + #[serde(skip_serializing_if = "Option::is_none")] + pub peak_memory_bytes: Option, } /// An object a run created, altered or recreated. @@ -1153,6 +1293,12 @@ 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, + /// The most memory the process held while the table loaded, in bytes; only under `--memory-per-table`. + #[serde(skip_serializing_if = "Option::is_none")] + pub peak_memory_bytes: Option, } /// 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..01e0d2bee 100644 --- a/crates/sqlload/src/loader/tests.rs +++ b/crates/sqlload/src/loader/tests.rs @@ -4,6 +4,8 @@ mod constant_batches; mod helpers; #[path = "tests/test_key_nulls.rs"] mod key_nulls; +#[path = "tests/test_lookback_start.rs"] +mod lookback_start; #[path = "tests/test_merged_files.rs"] mod merged_files; #[path = "tests/test_python_columns.rs"] @@ -19,5 +21,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_lookback_start.rs b/crates/sqlload/src/loader/tests/test_lookback_start.rs new file mode 100644 index 000000000..a472ea557 --- /dev/null +++ b/crates/sqlload/src/loader/tests/test_lookback_start.rs @@ -0,0 +1,56 @@ +use crate::column_types::models::ColumnType; +use crate::loader::_helpers::table_plan::lookback_start; +use crate::loader::tests::test_types::LookbackStartTestCase; + +#[test] +fn given_an_integer_lookback_when_counting_back_from_the_destination_maximum_then_the_start_stays_within_the_cursor_type() +-> Result<(), String> { + let test_cases = [ + LookbackStartTestCase { + description: "counts back the lookback's units", + maximum: "2026", + units: 1, + cursor_type: ColumnType::Int, + expected_start: Ok("2025"), + }, + LookbackStartTestCase { + description: "a smallint maximum at its lowest value stays there instead of leaving the type", + maximum: "-32768", + units: 1, + cursor_type: ColumnType::SmallInt, + expected_start: Ok("-32768"), + }, + LookbackStartTestCase { + description: "a smallint maximum near its lowest value stops at it", + maximum: "-32760", + units: 100, + cursor_type: ColumnType::SmallInt, + expected_start: Ok("-32768"), + }, + LookbackStartTestCase { + description: "an int maximum at its lowest value stays there", + maximum: "-9223372036854775808", + units: 5, + cursor_type: ColumnType::Int, + expected_start: Ok("-9223372036854775808"), + }, + LookbackStartTestCase { + description: "a maximum that isn't an integer can't be counted back from", + maximum: "2026.5", + units: 1, + cursor_type: ColumnType::Int, + expected_start: Err( + "cursor maximum \"2026.5\" is not an integer, so lookback = 1 can't count back from it", + ), + }, + ]; + for test_case in test_cases { + assert_eq!( + lookback_start(test_case.maximum, test_case.units, &test_case.cursor_type).as_deref(), + test_case.expected_start.map_err(str::to_owned).as_deref(), + "{}", + test_case.description + ); + } + Ok(()) +} 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_render_tally.rs b/crates/sqlload/src/loader/tests/test_render_tally.rs index 0769ae4a3..528892b18 100644 --- a/crates/sqlload/src/loader/tests/test_render_tally.rs +++ b/crates/sqlload/src/loader/tests/test_render_tally.rs @@ -60,12 +60,14 @@ fn given_a_load_when_rendering_its_tally_then_lines_align_and_the_last_counts_th retryable: false, elapsed: Duration::from_millis(*elapsed_ms), skipped: false, + peak_memory_bytes: None, }) .collect(); let (output, _) = render_tally( &applied, &outcomes, Duration::from_millis(test_case.elapsed_ms), + None, ); assert_eq!( output.join("\n"), 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..3828512a6 100644 --- a/crates/sqlload/src/loader/tests/test_types.rs +++ b/crates/sqlload/src/loader/tests/test_types.rs @@ -62,6 +62,37 @@ 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, +} + +pub(crate) struct LookbackStartTestCase { + pub(crate) description: &'static str, + pub(crate) maximum: &'static str, + pub(crate) units: i64, + pub(crate) cursor_type: crate::column_types::models::ColumnType, + pub(crate) expected_start: Result<&'static str, &'static str>, +} + /// A plan item: its name, its action, and each detail's label and text. pub(crate) type ItemSpec = ( &'static str, @@ -130,6 +161,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/number_text/constants.rs b/crates/sqlload/src/number_text/constants.rs new file mode 100644 index 000000000..3a576e2bd --- /dev/null +++ b/crates/sqlload/src/number_text/constants.rs @@ -0,0 +1,7 @@ +/// The step between binary memory units, and the units memory is shown in. +pub const BINARY_UNIT: u64 = 1024; +pub const MEMORY_UNITS: [&str; 4] = ["bytes", "KiB", "MiB", "GiB"]; + +/// Memory is shown to a tenth of its unit, and as a whole number from 100 of it (1,000 tenths) up. +pub const TENTHS: u64 = 10; +pub const WHOLE_FROM_TENTHS: u64 = 1_000; diff --git a/crates/sqlload/src/number_text/main/memory_text.rs b/crates/sqlload/src/number_text/main/memory_text.rs new file mode 100644 index 000000000..1a42c1130 --- /dev/null +++ b/crates/sqlload/src/number_text/main/memory_text.rs @@ -0,0 +1,31 @@ +//! Memory sizes in binary units. + +use crate::number_text::constants::{BINARY_UNIT, MEMORY_UNITS, TENTHS, WHOLE_FROM_TENTHS}; + +/// `bytes` in binary units to a tenth, as `512 MiB`, `1.9 GiB` or `40 KiB`; from 100 of a unit up, whole. +pub 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] + ), + } +} diff --git a/crates/sqlload/src/number_text/main/mod.rs b/crates/sqlload/src/number_text/main/mod.rs index 59890a8e2..b58bfd878 100644 --- a/crates/sqlload/src/number_text/main/mod.rs +++ b/crates/sqlload/src/number_text/main/mod.rs @@ -1 +1,2 @@ +pub mod memory_text; pub mod thousands; diff --git a/crates/sqlload/src/number_text/mod.rs b/crates/sqlload/src/number_text/mod.rs index 9b50e2b0a..6a04fe2e8 100644 --- a/crates/sqlload/src/number_text/mod.rs +++ b/crates/sqlload/src/number_text/mod.rs @@ -1,3 +1,4 @@ //! Numbers as text for people, as `plan`, `load` and the publish checks write them. +pub mod constants; pub mod main; 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..43304f360 --- /dev/null +++ b/crates/sqlload/src/odbc_source/_helpers/memory.rs @@ -0,0 +1,228 @@ +//! 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::number_text::constants::BINARY_UNIT; +use crate::number_text::main::memory_text::memory_text; +use crate::odbc_source::constants::{ + 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, + NAMED_LONG_COLUMNS, PROC_CGROUP, PROC_MEMINFO, UNLIMITED, +}; +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); + } + } +} + +/// `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::
--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/odbc.mdx b/website/src/content/docs/docs/sources/odbc.mdx index ff640ebf0..25b15d48c 100644 --- a/website/src/content/docs/docs/sources/odbc.mdx +++ b/website/src/content/docs/docs/sources/odbc.mdx @@ -1,6 +1,6 @@ --- title: "ODBC" -description: "Load tables and query results from any database with an ODBC driver: connections, choosing tables, query tables, value sizes and types." +description: "Load tables and query results from any database with an ODBC driver: connections, choosing tables, query tables, value sizes, memory and types." --- From sqlload 0.27.0, an `odbc` source loads tables, views and query @@ -23,7 +23,6 @@ connection = "erp" description = "Tables of the ERP database." include = ["sales.*"] exclude = ["sales.tmp_*"] -max_text_size = 8000 write_strategy = "replace" [sources.erp.tables."sales.items"] @@ -278,15 +277,105 @@ cursor outside it, so the query needs no placeholder. Write the query in the dat ## Value sizes Drivers read values into buffers of a fixed size per column, from the length the driver reports. -`max_text_size` (bytes) caps the buffer of every text column, and `max_binary_size` that of every -binary column, so set them when the driver reports a huge length for unbounded columns such as -`TEXT` or `VARCHAR(MAX)`, or none at all, which the read then refuses, naming the setting. A value -longer than its buffer, even in a `VARCHAR(10000)` column capped at 8000 bytes, fails the read, -naming the setting to raise; it's never cut short. +From 0.52.0, a text or binary column whose buffer would be over `long_column_threshold` (4 KiB when +omitted), or that the driver gives no length for, such as `TEXT`, `LONGTEXT` or `VARCHAR(MAX)`, is a +long column. A table with a long column is read value by value: each value at its own size, in +pieces for a long one, and never cut. A table without one is read in bulk, many rows per fetch. See +[Memory](#memory). + +`max_text_size` (bytes) caps every text column, and `max_binary_size` every binary column. They are +optional. A capped column counts as long only when the cap is over the threshold. A value longer than +its cap, even in a `VARCHAR(10000)` column capped at 8000 bytes, fails the read, naming the setting +to raise; it's never cut short. In a long column, the read stops as soon as the driver reports the +value's length or the pieces read pass the cap, so a 2 GB value over the cap isn't read first. Text +is compared as UTF-8 bytes, on Windows too, where drivers send it as UTF-16. From 0.47.0, the driver's own notes on the buffer size it chose for each column are no longer logged; set `RUST_LOG=info` to see them again. +## Memory + +From 0.52.0, each ODBC table has a fetch buffer: the most the rows of one fetch hold. It's 512 MiB +per table by default. On a smaller machine, it's half the memory shared between the tables loading +at once ([`concurrency`](/docs/configuration/)), when that is less. The memory is the container's +limit (the cgroup's `memory.max`) when SQLLoad runs in one, and otherwise the machine's memory. The +fetch buffer is never raised above 512 MiB unless the source sets `max_batch_memory`: + +```toml +[sources.shop] +kind = "odbc" +connection = "shop" +max_batch_memory = "2GB" +``` + +The fetch buffer bounds the source's fetches only, not the load's whole memory: the batches made +from them, the destination and the commits take memory besides. To measure what a load really +takes, read `peak_memory_bytes` in the [JSON report](/docs/cli/load/#results-as-json), or run once +with [`--memory-per-table`](/docs/cli/load/#peak-memory) on Linux; see +[Finding a source that uses too much memory](/docs/troubleshooting/#finding-a-source-that-uses-too-much-memory). + +`plan` shows how each ODBC table fetches (`fetch`) and its fetch buffer, and why (`buffer`): + +```console +$ sqlload plan +… +Loads (2) sqlload load +├── raw.shop__contacts replace +│ ├── source odbc shop contacts +│ ├── route read as Arrow -> DuckDB takes Arrow +│ ├── fetch up to 512 MiB or 65,536 rows per fetch; 265 long columns (c001, c002, c003, …), so all 266 columns are read value by value +│ └── buffer 512 MiB fetch buffer per table (default; 6 GiB container limit, concurrency 4) +└── raw.shop__orders replace + ├── source odbc shop orders + ├── route read as Arrow -> DuckDB takes Arrow + ├── fetch ~13.6 MiB per fetch, 65,536 rows + └── buffer 512 MiB fetch buffer per table (default; 6 GiB container limit, concurrency 4) +``` + +- **In bulk** (`orders`): a fetch reads as many rows as its buffers hold within the fetch buffer, + at most 65,536. The `fetch` line shows the buffers of one fetch. +- **Value by value** (`contacts`): the driver returns one row at a time, and SQLLoad reads each of + its values. A fetch ends at 65,536 rows, or once the values it has read reach the fetch buffer. A + value larger than the whole fetch buffer still loads; its fetch ends after its row, and SQLLoad + logs it: + + ```text + … WARN sqlload::odbc_source::_helpers::values: a value is larger than the table's fetch buffer (max_batch_memory); it loads, and its fetch ends after its row table="notes" column=body bytes=40000000 budget=16000000 + ``` + +A table with a long column fetches one row at a time because reading values with `SQLGetData` +from a block of rows breaks on common drivers: on MariaDB Connector/ODBC 3.1, every row after the +first returns `SQL_NO_DATA`, and on Microsoft ODBC Driver 18 for SQL Server, `SQLSetPos` fails on +the forward-only cursor. Tables without a long column keep fetching in bulk. + +The fetch buffer sets only how many rows a fetch reads. Commits stay as `batch_size` sets them, so +a smaller fetch buffer means more fetches, not different data. + +When the fetch buffers of the tables that may load at once together pass half the memory, `plan` +warns, and `load` checks again when it starts, since it may run on another machine. The warning +counts the fetch buffers only. Here the source sets `max_batch_memory = "4GB"` on a machine with a +6 GiB container limit: + +```console +$ sqlload plan --concurrency 2 +… +Warnings (1) +└── the ODBC fetch buffers of 2 tables loading at once may hold 3.7 GiB together, more than half of this machine's 6 GiB memory limit; this counts only those buffers, not the memory the batches, the destination and the commits take besides; lower max_batch_memory on source shop, or lower concurrency (now 2) +``` + +| Key | Default | Meaning | +|-----|---------|---------| +| `max_batch_memory` | 512 MiB, or less on a small machine | The fetch buffer of each table: the most the rows of one fetch hold, such as `"2GB"` or `"256MiB"`. It doesn't bound the load's whole memory. | +| `long_column_threshold` | `"4KiB"` | The buffer per value over which a text or binary column is read value by value. | + +Value by value is slower than bulk: SQLLoad asks the driver for each value with its own call +(`SQLGetData`), as pyodbc reads every value +([pyodbc 5.2.0, `getdata.cpp`](https://github.com/mkleehammer/pyodbc/blob/5.2.0/src/getdata.cpp)). +It saves memory when most values are much shorter than their column allows, as in wide tables of +short notes. When the values really are long, they take the memory either way. A table with one +long column is read value by value in all its columns; raise `long_column_threshold` above the +column's size to read it in bulk, with buffers of that size. + ## Types SQLLoad reads each column's type as the driver describes it: diff --git a/website/src/content/docs/docs/sources/postgres-cdc.mdx b/website/src/content/docs/docs/sources/postgres-cdc.mdx index 7f63675e0..5ffd16d00 100644 --- a/website/src/content/docs/docs/sources/postgres-cdc.mdx +++ b/website/src/content/docs/docs/sources/postgres-cdc.mdx @@ -154,11 +154,12 @@ from [step 4](/docs/guides/postgres-cdc/#4-load-the-changes): $ sqlload load ✓ Plan complete 2 selected, 2 to load (sqlload plan shows the tree) -loaded lake/app/customers/ 50 rows 0.2s -loaded lake/app/orders/ 1,000 rows 0.2s +loaded lake/app/customers/ 50 rows 0.3s +loaded lake/app/orders/ 1,000 rows 0.3s -✓ Completed 2 loaded, 1,050 rows, 0.2s -└── slowest lake/app/customers/ 0.2s +✓ Completed 2 loaded, 1,050 rows, 0.3s +├── slowest lake/app/customers/ 0.3s +└── peak memory 66.6 MiB $ sqlload load ✓ Plan complete 2 selected, 2 to load (sqlload plan shows the tree) @@ -166,7 +167,8 @@ loaded lake/app/customers/ 0 rows 0.3s loaded lake/app/orders/ 3 rows 0.3s ✓ Completed 2 loaded, 3 rows, 0.3s -└── slowest lake/app/customers/ 0.3s +├── slowest lake/app/customers/ 0.3s +└── peak memory 65.3 MiB ``` Then read the folder with DuckDB: diff --git a/website/src/content/docs/docs/sources/python.mdx b/website/src/content/docs/docs/sources/python.mdx index 4e7812f73..c648de2c1 100644 --- a/website/src/content/docs/docs/sources/python.mdx +++ b/website/src/content/docs/docs/sources/python.mdx @@ -49,7 +49,7 @@ Unknown settings raise `TypeError` when the module is imported. | Setting | On | Meaning | |---------|----|---------| -| `write_strategy`, `unique_key`, `on_duplicate_keys`, `null_keys`, `dedup_sort`, `cursor_column`, `cursor_start`, `lookback`, `append_cursor_inclusive`, `batch_size` | source or table | As in [`sqlload.toml`](/docs/configuration/#table-settings), such as `dedup_sort={"updated_at": "desc"}`; `on_duplicate_keys` and `dedup_sort` from 0.39.0, `null_keys` from 0.51.0. Python tables don't take `lookahead` or `lookback_from`. Every table needs a `write_strategy`, from its source or its own decorator, and a `merge` table names its `unique_key` columns: a Python source reports no primary key, so `"{primary_key}"` is refused. | +| `write_strategy`, `unique_key`, `on_duplicate_keys`, `null_keys`, `dedup_sort`, `cursor_column`, `cursor_start`, `lookback`, `append_cursor_inclusive`, `batch_size` | source or table | As in [`sqlload.toml`](/docs/configuration/#table-settings), such as `dedup_sort={"updated_at": "desc"}`; `on_duplicate_keys` and `dedup_sort` from 0.39.0, `null_keys` from 0.51.0, and `lookback=1` on an integer cursor from 0.52.0. Python tables don't take `lookahead` or `lookback_from`. Every table needs a `write_strategy`, from its source or its own decorator, and a `merge` table names its `unique_key` columns: a Python source reports no primary key, so `"{primary_key}"` is refused. | | `contract`, `on_schema_change`, `auto_widen`, `type_change_rebuild` | source or table | See [Contracts](/docs/concepts/contracts/). | | `checks` | source or table | Checks on the staged rows, such as `{"unique": ["order_id"]}`. See [Checks before publish](/docs/concepts/checks/). | | `constant_columns` | source or table | From 0.37.0. Columns with the same value in every row, such as `{"tenant_id": "acme"}` or `{"shard": {"value": "7", "type": "int"}}`; a table's replace its source's. A yielded field of the same name, ignoring case, fails the table before anything is written. See [Several tenants in one table](/docs/concepts/write-strategies/#several-tenants-in-one-table). | @@ -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 @@ -268,7 +308,7 @@ A table function receives a `sl.Context`: | Attribute | Meaning | |-----------|---------| -| `ctx.cursor_value` | Where to read from: the destination's maximum minus `lookback`, `cursor_start`, or `--start-cursor`. `None` reads everything. | +| `ctx.cursor_value` | Where to read from: the destination's maximum minus `lookback`, `cursor_start`, or `--start-cursor` (from 0.52.0, or the typed flag of the cursor's kind). `None` reads everything. | | `ctx.cursor_end` | Where the window ends, with `batch_size` or `--end-cursor`. | | `ctx.cursor_column`, `ctx.cursor_inclusive` | The cursor column, and whether the read includes `cursor_value`. `False` when `append` reads after the destination's maximum. | | `ctx.is_reload` | `True` under `--reload`. | @@ -295,8 +335,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 +367,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 +409,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 +696,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. | diff --git a/website/src/content/docs/docs/targets/files.mdx b/website/src/content/docs/docs/targets/files.mdx index 9648e48b9..630840e72 100644 --- a/website/src/content/docs/docs/targets/files.mdx +++ b/website/src/content/docs/docs/targets/files.mdx @@ -35,6 +35,7 @@ Loads (1) sqlload load └── app__orders delete_insert (cursor: updated_at) ├── source postgres app public.orders ├── reads updated_at: first run, everything + ├── columns 5 ├── files out/app/orders/ parquet (zstd) partitions dt={cursor:date}/ ├── schema first files, 7 columns ├── note each partition written or in the window gets its new files first, then its older files are deleted; not atomic, so a reader in between sees both @@ -47,7 +48,8 @@ $ sqlload load loaded out/app/orders/ 1,000 rows 0.1s ✓ Completed 1 loaded, 1,000 rows, 0.1s -└── slowest out/app/orders/ 0.1s +├── slowest out/app/orders/ 0.1s +└── peak memory 66.2 MiB ``` Files targets take `append`, `replace` and `delete_insert`. `merge` needs a table format that diff --git a/website/src/content/docs/docs/troubleshooting.mdx b/website/src/content/docs/docs/troubleshooting.mdx index 4c5d4f5fc..dba7df491 100644 --- a/website/src/content/docs/docs/troubleshooting.mdx +++ b/website/src/content/docs/docs/troubleshooting.mdx @@ -17,6 +17,8 @@ common ones; each message is copied from SQLLoad. | `public.events (kind = "postgres"): no write_strategy is set on the table or its source, and SQLLoad doesn't choose one` | From 0.46.0, every table needs a `write_strategy`. Add the line the error shows under the source, for all its tables, or under the table's entry. See [Write strategies](/docs/concepts/write-strategies/). | | `write_strategy = "merge" needs a unique_key: the key columns, such as unique_key = ["id"], or unique_key = "{primary_key}" ...` | From 0.46.0, a `merge` table names its key. Set `unique_key` to the columns, or to `"{primary_key}"` for the primary key the source reports. | | `items (kind = "odbc"): unique_key = "{primary_key}", but the source reports no primary key for the table. ...` | Name the key columns. An ODBC table whose error adds that `INFORMATION_SCHEMA` lists a primary key has one its driver didn't report; see [MariaDB Connector/ODBC 3.1](/docs/sources/odbc/#mariadb-connectorodbc-31). | +| `the ODBC fetch buffers of 2 tables loading at once may hold 3.7 GiB together, more than half of this machine's 6 GiB memory limit; ...` | A warning from 0.52.0. It counts the fetch buffers only. Lower `max_batch_memory` on the named source, or `concurrency`. See [Memory](/docs/sources/odbc/#memory). | +| `column body holds a value of 9000 bytes, longer than max_text_size = 8000; ...` | Raise `max_text_size` on the ODBC source, or remove it: from 0.52.0 a long column is read value by value without it. See [Value sizes](/docs/sources/odbc/#value-sizes). | | `write_strategy = "replace" reloads the whole table, so it takes no cursor_column` | Remove `cursor_column`, or choose `merge`, `delete_insert` or `append`. | | `lookback re-inserts the window on every run with append; use delete_insert (no key needed) or merge` | Change the strategy, or drop `lookback`. | | `merge needs a table format to update rows in place; a files target takes append, replace or delete_insert` | Use one of those strategies on a files target. | @@ -61,6 +63,47 @@ such as `copy_to_internal_stage`, `hmac`-masked files, or `load_method = { type files into Snowflake, an external stage over the files' bucket lets Snowflake read them in place. See [Load methods](/docs/targets/). +For an ODBC table, `plan`'s `fetch` line says when a table is read value by value, which is slower +than bulk. See [Memory](/docs/sources/odbc/#memory). + +## Finding a source that uses too much memory + +From 0.52.0, every load ends with its peak memory: `peak memory` in the summary, and +`peak_memory_bytes` in the [JSON report](/docs/cli/load/#results-as-json) and the `run_finished` +event. It's the peak for the whole run, all tables together, ODBC drivers, DuckDB and Python +sources included. Under the [Python package](/docs/guides/python/), it includes the Python +interpreter. + +To find the table behind it, run once on Linux with +[`--memory-per-table`](/docs/cli/load/#peak-memory). It loads one table at a time and gives each +table's peak: + +```console +$ sqlload load --memory-per-table +… +loaded raw.customers 500 rows 0.0s peak memory 79.8 MiB +loaded raw.orders 2,000,000 rows 8.5s peak memory 245 MiB + +✓ Completed 2 loaded, 2,000,500 rows, 8.6s +├── slowest raw.orders 8.5s +└── peak memory 245 MiB +``` + +A table's figure includes what the process already held when it started, so compare the tables of +one run. Then, for the table it names: + +- Lower [`concurrency`](/docs/configuration/), or pass `--concurrency`, as tables loading at once + add up. +- Set a smaller [`batch_size`](/docs/concepts/cursors/#batches) on a cursor table, so it reads and + commits smaller windows, one at a time. +- On an [ODBC source](/docs/sources/odbc/#memory), lower `max_batch_memory`, its fetch buffer. + When wide text or binary columns hold short values, lower `long_column_threshold` so they're read + value by value, which holds each value's real size instead of the column's whole width, but is + slower. + +Run it again with `--memory-per-table` to check. On macOS and Windows, load the table on its own +(`sqlload load TABLE`) and compare its run's peak with the others'. + ## A large first load A cursor table without an ordered read commits once for the whole read. For a large first load,