feat: named transforms, ODBC fetch buffer, peak memory, typed cursor flags, plan-time path checks and streamed --follow

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 <function> 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.

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"
 

crates/sqlload-python/src/crossing/_helpers/parser.rs

@@ -6,18 +6,25 @@ 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;
 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,31 +263,26 @@ 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<bool, CrossError> {
+        let transform = match serde_json::from_str::<HttpJob>(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<String> = None;
         run_http_job(
             definition,
             &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<bool, String> {
+        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<RecordBatch> = 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<RecordBatch>,
+    ) -> Result<bool, String> {
+        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());

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);

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,
+    }
+}

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;

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<RecordBatch>, SchemaRef),
+    (options, flatten): (CrossingOptions, Vec<String>),
+) -> Result<(TableBatches, Py<PyAny>), String> {
+    let context = serde_json::to_string(&call.context).map_err(|error| error.to_string())?;
+    Python::attach(
+        |py| -> PyResult<Result<(TableBatches, Py<PyAny>), String>> {
+            let table = Table::try_new(rows, schema)
+                .map_err(|error| PyValueError::new_err(error.to_string()))?
+                .into_pyarrow(py)?;
+            let (returned, problem): (Option<Bound<'_, PyAny>>, Option<String>) = 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())?
+}

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<String>,
     pub connections: serde_json::Map<String, serde_json::Value>,
     /// The table's cursor column, which projected JSON rows must keep.