From df6063d2dc1cdcff5d92df12d08299de1c782381 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 15 Oct 2026 09:14:10 +0100 Subject: [PATCH 01/12] refactor: move compile timings, built-in audits and code identity to native --- Cargo.lock | 1 + .../src/audits/constants.rs | 19 +++ crates/sqlbuild-attachments/src/audits/mod.rs | 2 +- crates/sqlbuild-cache/Cargo.toml | 1 + .../src/digest/_helpers/distributions.rs | 124 ++++++++++++++++++ .../sqlbuild-cache/src/digest/_helpers/mod.rs | 1 + crates/sqlbuild-cache/src/digest/constants.rs | 13 ++ .../src/digest/main/installed_code_facts.rs | 44 +++++++ crates/sqlbuild-cache/src/digest/main/mod.rs | 1 + crates/sqlbuild-cache/src/digest/tests.rs | 2 + .../digest/tests/test_installed_code_facts.rs | 88 +++++++++++++ .../src/digest/tests/test_types.rs | 9 ++ .../src/compile_timings/constants.rs | 14 +- .../main/compile_timing_milliseconds.rs | 19 +++ .../src/compile_timings/main/mod.rs | 1 + .../sqlbuild-core/src/compile_timings/mod.rs | 8 ++ .../src/compile_timings/models.rs | 10 ++ .../src/compile_timings/tests.rs | 4 + .../tests/test_compile_timing_milliseconds.rs | 37 ++++++ .../src/compile_timings/tests/test_types.rs | 7 + crates/sqlbuild-core/src/lib.rs | 1 + .../analysis_session/model_analyses.rs | 15 ++- .../bindings/_helpers/cache/native_store.rs | 78 +++++++++-- .../_helpers/cache/sql_test_scan_store.rs | 27 ++-- .../compile_command/artifact_batch.rs | 18 +-- .../_helpers/compile_command/target_writer.rs | 11 +- .../_helpers/compile_driver/analysis_cache.rs | 12 +- .../compile_driver/assembly_analysis.rs | 16 +-- .../_helpers/compile_driver/command.rs | 27 ++-- .../compile_driver/compile_timings.rs | 98 ++++++++++++++ .../bindings/_helpers/compile_driver/mod.rs | 1 + .../_helpers/compile_driver/project.rs | 11 +- .../_helpers/compiled_project/attachment.rs | 55 ++++++-- .../src/bindings/_helpers/rules/rows.rs | 6 +- .../_helpers/sql_test_glue/plan_driver.rs | 41 ++---- fensu.toml | 5 + python-boundary.toml | 22 +--- src/sqlbuild/_native.pyi | 3 +- .../compiler/auditing/main/_builtins.py | 83 ------------ .../compiler/frontier/_helpers/__init__.py | 0 .../frontier/_helpers/code_identity.py | 62 --------- .../compiler/frontier/main/__init__.py | 1 - .../frontier/main/compiled_code_identity.py | 10 -- src/sqlbuild/compiler/profiling/__init__.py | 1 - .../compiler/profiling/_helpers/__init__.py | 1 - .../compiler/profiling/_helpers/timing.py | 23 ---- .../compiler/profiling/classes/__init__.py | 1 - .../compiler/profiling/classes/context.py | 14 -- src/sqlbuild/compiler/profiling/constants.py | 28 ---- .../compiler/profiling/main/__init__.py | 1 - .../compiler/profiling/main/_collect.py | 20 --- .../compiler/profiling/main/_metric.py | 13 -- .../compiler/profiling/main/_record.py | 13 -- .../compiler/profiling/main/_record_cpu.py | 13 -- src/sqlbuild/compiler/profiling/models.py | 31 ----- .../_helpers/engine/custom_dispatch.py | 3 +- .../_helpers/engine/custom_rules.py | 4 +- .../cli/commands/main/compile/helpers.py | 5 +- .../rule_engine/_helpers/run/helpers.py | 5 +- 59 files changed, 685 insertions(+), 499 deletions(-) create mode 100644 crates/sqlbuild-cache/src/digest/_helpers/distributions.rs create mode 100644 crates/sqlbuild-cache/src/digest/main/installed_code_facts.rs create mode 100644 crates/sqlbuild-cache/src/digest/tests/test_installed_code_facts.rs rename src/sqlbuild/compiler/profiling/types.py => crates/sqlbuild-core/src/compile_timings/constants.rs (66%) create mode 100644 crates/sqlbuild-core/src/compile_timings/main/compile_timing_milliseconds.rs create mode 100644 crates/sqlbuild-core/src/compile_timings/main/mod.rs create mode 100644 crates/sqlbuild-core/src/compile_timings/mod.rs create mode 100644 crates/sqlbuild-core/src/compile_timings/models.rs create mode 100644 crates/sqlbuild-core/src/compile_timings/tests.rs create mode 100644 crates/sqlbuild-core/src/compile_timings/tests/test_compile_timing_milliseconds.rs create mode 100644 crates/sqlbuild-core/src/compile_timings/tests/test_types.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/compile_driver/compile_timings.rs delete mode 100644 src/sqlbuild/compiler/auditing/main/_builtins.py delete mode 100644 src/sqlbuild/compiler/frontier/_helpers/__init__.py delete mode 100644 src/sqlbuild/compiler/frontier/_helpers/code_identity.py delete mode 100644 src/sqlbuild/compiler/frontier/main/__init__.py delete mode 100644 src/sqlbuild/compiler/frontier/main/compiled_code_identity.py delete mode 100644 src/sqlbuild/compiler/profiling/__init__.py delete mode 100644 src/sqlbuild/compiler/profiling/_helpers/__init__.py delete mode 100644 src/sqlbuild/compiler/profiling/_helpers/timing.py delete mode 100644 src/sqlbuild/compiler/profiling/classes/__init__.py delete mode 100644 src/sqlbuild/compiler/profiling/classes/context.py delete mode 100644 src/sqlbuild/compiler/profiling/constants.py delete mode 100644 src/sqlbuild/compiler/profiling/main/__init__.py delete mode 100644 src/sqlbuild/compiler/profiling/main/_collect.py delete mode 100644 src/sqlbuild/compiler/profiling/main/_metric.py delete mode 100644 src/sqlbuild/compiler/profiling/main/_record.py delete mode 100644 src/sqlbuild/compiler/profiling/main/_record_cpu.py delete mode 100644 src/sqlbuild/compiler/profiling/models.py diff --git a/Cargo.lock b/Cargo.lock index 14631c2f4f..dbfac92b5c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1175,6 +1175,7 @@ version = "0.156.1" dependencies = [ "blake3", "rayon", + "serde_json", "sha2", "sqlbuild-core", "tempfile", diff --git a/crates/sqlbuild-attachments/src/audits/constants.rs b/crates/sqlbuild-attachments/src/audits/constants.rs index b0bf059d59..657b260741 100644 --- a/crates/sqlbuild-attachments/src/audits/constants.rs +++ b/crates/sqlbuild-attachments/src/audits/constants.rs @@ -13,3 +13,22 @@ pub(crate) const FUNCTION_RESOURCE: &str = "function"; pub(crate) const CROSS_MODEL_MINIMUM: usize = 2; /// The run scope Python falls back to without an instance or project value. pub(crate) const FALLBACK_RUN_SCOPE: &str = "delta_and_final"; +/// Built-in generic audit names and SQL bodies, which project definitions override by name. +pub const BUILTIN_AUDITS: [(&str, &str); 4] = [ + ( + "accepted_values", + "SELECT @column\nFROM @relation\nWHERE @column IS NOT NULL\n AND @column NOT IN (@'values')", + ), + ( + "not_null", + "SELECT @column\nFROM @relation\nWHERE @column IS NULL", + ), + ( + "relationships", + "SELECT @column\nFROM @relation\nWHERE @column IS NOT NULL\n AND @column NOT IN (\n SELECT @field\n FROM @to\n WHERE @field IS NOT NULL\n )", + ), + ( + "unique", + "SELECT @column, COUNT(*) AS duplicate_count\nFROM @relation\nWHERE @column IS NOT NULL\nGROUP BY @column\nHAVING COUNT(*) > 1", + ), +]; diff --git a/crates/sqlbuild-attachments/src/audits/mod.rs b/crates/sqlbuild-attachments/src/audits/mod.rs index 1e64f3d116..f01181f853 100644 --- a/crates/sqlbuild-attachments/src/audits/mod.rs +++ b/crates/sqlbuild-attachments/src/audits/mod.rs @@ -1,5 +1,5 @@ pub(crate) mod _helpers; -pub(crate) mod constants; +pub mod constants; pub mod main; pub mod models; diff --git a/crates/sqlbuild-cache/Cargo.toml b/crates/sqlbuild-cache/Cargo.toml index 0116e0d5e3..617208de4e 100644 --- a/crates/sqlbuild-cache/Cargo.toml +++ b/crates/sqlbuild-cache/Cargo.toml @@ -10,6 +10,7 @@ rust-version.workspace = true [dependencies] blake3.workspace = true rayon.workspace = true +serde_json.workspace = true sha2.workspace = true sqlbuild-core.workspace = true tempfile.workspace = true diff --git a/crates/sqlbuild-cache/src/digest/_helpers/distributions.rs b/crates/sqlbuild-cache/src/digest/_helpers/distributions.rs new file mode 100644 index 0000000000..67320b5f38 --- /dev/null +++ b/crates/sqlbuild-cache/src/digest/_helpers/distributions.rs @@ -0,0 +1,124 @@ +//! `importlib.metadata` distribution lookup over search paths, for code identity facts. + +use std::io::ErrorKind; +use std::path::{Path, PathBuf}; + +use crate::digest::constants::{METADATA_DIRECTORY_SUFFIXES, METADATA_FILES}; + +/// The first metadata folder of `name` along `search_paths`, as `importlib.metadata` finds it. +pub(crate) fn distribution_directory(search_paths: &[PathBuf], name: &str) -> Option { + let wanted: String = normalized(name); + for root in search_paths { + let directory: &Path = if root.as_os_str().is_empty() { + Path::new(".") + } else { + root + }; + if let Some(child) = children(directory) + .into_iter() + .find(|child| metadata_name(child).is_some_and(|found| normalized(&found) == wanted)) + { + return Some(root.join(child)); + } + } + None +} + +/// The `Version` header of a distribution's metadata, when it has one. +pub(crate) fn distribution_version(directory: &Path) -> Option { + for file in METADATA_FILES { + if let Some(text) = optional_text(&directory.join(file)) { + return header_value(&text, "version"); + } + } + None +} + +/// A file's bytes, or `None` where `importlib.metadata` treats the file as absent. +pub(crate) fn optional_bytes(path: &Path) -> Option> { + match std::fs::read(path) { + Ok(bytes) => Some(bytes), + Err(error) if absent(&error) => None, + Err(_) => None, + } +} + +/// Entry names of a search path folder; folders Python cannot list have none. +fn children(directory: &Path) -> Vec { + let entries: std::fs::ReadDir = match std::fs::read_dir(directory) { + Ok(entries) => entries, + Err(error) if absent(&error) => return Vec::new(), + Err(_) => return Vec::new(), + }; + let mut names: Vec = Vec::new(); + for entry in entries.flatten() { + names.push(entry.file_name().to_string_lossy().into_owned()); + } + names +} + +fn optional_text(path: &Path) -> Option { + optional_bytes(path).map(|bytes| String::from_utf8_lossy(&bytes).into_owned()) +} + +fn absent(error: &std::io::Error) -> bool { + matches!( + error.kind(), + ErrorKind::NotFound | ErrorKind::NotADirectory | ErrorKind::IsADirectory + ) +} + +/// The distribution name a metadata folder entry names, before its version. +fn metadata_name(child: &str) -> Option { + let lowered: String = child.to_lowercase(); + METADATA_DIRECTORY_SUFFIXES + .iter() + .any(|suffix| lowered.ends_with(suffix)) + .then(|| { + let stem: &str = lowered.rsplit_once('.').map_or("", |(stem, _)| stem); + stem.split('-').next().unwrap_or_default().to_owned() + }) +} + +/// `re.sub(r"[-_.]+", "-", name).lower().replace("-", "_")`. +fn normalized(name: &str) -> String { + let mut text: String = String::with_capacity(name.len()); + let mut separator: bool = false; + for character in name.chars() { + if matches!(character, '-' | '_' | '.') { + if !separator { + text.push('_'); + } + separator = true; + } else { + text.extend(character.to_lowercase()); + separator = false; + } + } + text +} + +/// The first value of a case-insensitive header in an RFC 822 style header block. +fn header_value(text: &str, name: &str) -> Option { + let mut value: Option = None; + for line in text.lines() { + if line.is_empty() { + break; + } + if line.starts_with([' ', '\t']) { + if let Some(current) = value.as_mut() { + current.push_str(line); + } + continue; + } + if value.is_some() { + break; + } + if let Some((header, rest)) = line.split_once(':') + && header.eq_ignore_ascii_case(name) + { + value = Some(rest.trim_start().to_owned()); + } + } + value.map(|text| text.trim_end().to_owned()) +} diff --git a/crates/sqlbuild-cache/src/digest/_helpers/mod.rs b/crates/sqlbuild-cache/src/digest/_helpers/mod.rs index 2ebc7ff525..c32b54c6a3 100644 --- a/crates/sqlbuild-cache/src/digest/_helpers/mod.rs +++ b/crates/sqlbuild-cache/src/digest/_helpers/mod.rs @@ -1,2 +1,3 @@ pub(crate) mod code_files; +pub(crate) mod distributions; pub(crate) mod walk; diff --git a/crates/sqlbuild-cache/src/digest/constants.rs b/crates/sqlbuild-cache/src/digest/constants.rs index 8494ec4831..3f14686585 100644 --- a/crates/sqlbuild-cache/src/digest/constants.rs +++ b/crates/sqlbuild-cache/src/digest/constants.rs @@ -35,3 +35,16 @@ pub const PRESENCE_TAG: &[u8] = b"presence"; pub(crate) const CODE_SOURCE_SUFFIXES: &[&str] = &["py", "so", "pyd"]; /// The bytecode folder an editable install's code walk skips. pub(crate) const BYTECODE_CACHE_DIRECTORY: &str = "__pycache__"; + +/// Distributions whose installed versions identify the code behind cached results. +pub const CODE_IDENTITY_DISTRIBUTIONS: [&str; 3] = ["sqlbuild", "polyglot-sql-chio", "PyYAML"]; +/// The distribution an editable install's `direct_url.json` describes. +pub(crate) const SQLBUILD_DISTRIBUTION: &str = "sqlbuild"; +/// The version reported for a distribution no search path holds. +pub(crate) const UNKNOWN_VERSION: &str = "unknown"; +/// Metadata folder suffixes `importlib.metadata` searches, lowercased. +pub(crate) const METADATA_DIRECTORY_SUFFIXES: [&str; 2] = [".dist-info", ".egg-info"]; +/// Metadata files read for a distribution's version, in order. +pub(crate) const METADATA_FILES: [&str; 2] = ["METADATA", "PKG-INFO"]; +/// The installer's record of where a distribution was installed from. +pub(crate) const DIRECT_URL_FILE: &str = "direct_url.json"; diff --git a/crates/sqlbuild-cache/src/digest/main/installed_code_facts.rs b/crates/sqlbuild-cache/src/digest/main/installed_code_facts.rs new file mode 100644 index 0000000000..6157b0856a --- /dev/null +++ b/crates/sqlbuild-cache/src/digest/main/installed_code_facts.rs @@ -0,0 +1,44 @@ +//! The installed distribution versions and editable flag behind the installed code identity. + +use std::path::PathBuf; + +use crate::digest::_helpers::distributions::{ + distribution_directory, distribution_version, optional_bytes, +}; +use crate::digest::constants::{ + CODE_IDENTITY_DISTRIBUTIONS, DIRECT_URL_FILE, SQLBUILD_DISTRIBUTION, UNKNOWN_VERSION, +}; + +/// Each identity distribution's version along `search_paths`, and whether SQLBuild is editable. +#[must_use] +pub fn installed_code_facts(search_paths: &[PathBuf]) -> (Vec, bool) { + let versions: Vec = CODE_IDENTITY_DISTRIBUTIONS + .iter() + .map(|name| { + distribution_directory(search_paths, name) + .and_then(|directory| distribution_version(&directory)) + .unwrap_or_else(|| UNKNOWN_VERSION.to_owned()) + }) + .collect(); + let editable: bool = + distribution_directory(search_paths, SQLBUILD_DISTRIBUTION).is_none_or(|directory| { + optional_bytes(&directory.join(DIRECT_URL_FILE)) + .is_some_and(|bytes| editable_direct_url(&bytes)) + }); + (versions, editable) +} + +/// Whether `direct_url.json` marks an editable install; unreadable payloads count as editable. +fn editable_direct_url(bytes: &[u8]) -> bool { + let Ok(payload) = serde_json::from_slice::(bytes) else { + return true; + }; + let Some(payload) = payload.as_object() else { + return true; + }; + payload + .get("dir_info") + .and_then(serde_json::Value::as_object) + .and_then(|dir_info| dir_info.get("editable")) + .is_some_and(|editable| editable == &serde_json::Value::Bool(true)) +} diff --git a/crates/sqlbuild-cache/src/digest/main/mod.rs b/crates/sqlbuild-cache/src/digest/main/mod.rs index fb6e1e2b02..3d1fb78930 100644 --- a/crates/sqlbuild-cache/src/digest/main/mod.rs +++ b/crates/sqlbuild-cache/src/digest/main/mod.rs @@ -4,5 +4,6 @@ pub mod content_digest; pub mod digest_files; pub mod fingerprint_project_files; pub mod hex_digest; +pub mod installed_code_facts; pub mod installed_code_identity; pub mod sha256_hex; diff --git a/crates/sqlbuild-cache/src/digest/tests.rs b/crates/sqlbuild-cache/src/digest/tests.rs index 329720ef78..c607d56b8b 100644 --- a/crates/sqlbuild-cache/src/digest/tests.rs +++ b/crates/sqlbuild-cache/src/digest/tests.rs @@ -8,5 +8,7 @@ mod digest_files; mod fingerprint_project_files; #[path = "tests/helpers.rs"] mod helpers; +#[path = "tests/test_installed_code_facts.rs"] +mod installed_code_facts; #[path = "tests/test_types.rs"] mod test_types; diff --git a/crates/sqlbuild-cache/src/digest/tests/test_installed_code_facts.rs b/crates/sqlbuild-cache/src/digest/tests/test_installed_code_facts.rs new file mode 100644 index 0000000000..52c1531656 --- /dev/null +++ b/crates/sqlbuild-cache/src/digest/tests/test_installed_code_facts.rs @@ -0,0 +1,88 @@ +use std::path::PathBuf; + +use crate::digest::main::installed_code_facts::installed_code_facts; +use crate::digest::tests::helpers::write_file; +use crate::digest::tests::test_types::InstalledCodeFactsTestCase; + +#[test] +fn given_search_paths_when_reading_installed_facts_then_importlib_metadata_is_reproduced() { + let test_cases = [ + InstalledCodeFactsTestCase { + description: "nothing installed is unknown and editable", + files: &[(0, "other.txt", "")], + expected_versions: ["unknown", "unknown", "unknown"], + expected_editable: true, + }, + InstalledCodeFactsTestCase { + description: "the first search path wins and names normalise", + files: &[ + ( + 0, + "sqlbuild-1.2.0.dist-info/METADATA", + "Name: sqlbuild\nVersion: 1.2.0\n\nbody", + ), + (1, "sqlbuild-9.0.0.dist-info/METADATA", "Version: 9.0.0\n"), + ( + 1, + "polyglot_sql_chio-0.4.dist-info/METADATA", + "version: 0.4 \n", + ), + ( + 1, + "PyYAML-6.0.2.dist-info/METADATA", + "Metadata-Version: 2.1\nVersion: 6.0.2\n", + ), + ], + expected_versions: ["1.2.0", "0.4", "6.0.2"], + expected_editable: false, + }, + InstalledCodeFactsTestCase { + description: "an editable direct URL marks the install editable", + files: &[ + (0, "sqlbuild-1.0.dist-info/METADATA", "Version: 1.0\n"), + ( + 0, + "sqlbuild-1.0.dist-info/direct_url.json", + "{\"url\": \"file:///src\", \"dir_info\": {\"editable\": true}}", + ), + ], + expected_versions: ["1.0", "unknown", "unknown"], + expected_editable: true, + }, + InstalledCodeFactsTestCase { + description: "an unreadable direct URL counts as editable", + files: &[ + (0, "sqlbuild-1.0.dist-info/METADATA", "Version: 1.0\n"), + (0, "sqlbuild-1.0.dist-info/direct_url.json", "not json"), + ], + expected_versions: ["1.0", "unknown", "unknown"], + expected_editable: true, + }, + ]; + for test_case in test_cases { + let directories = [ + tempfile::tempdir().expect("a temporary folder"), + tempfile::tempdir().expect("a temporary folder"), + ]; + for (index, relative_path, contents) in test_case.files { + write_file(directories[*index].path(), relative_path, contents); + } + let search_paths: Vec = directories + .iter() + .map(|directory| directory.path().to_path_buf()) + .collect(); + + let (versions, editable) = installed_code_facts(&search_paths); + + assert_eq!( + versions, test_case.expected_versions, + "{}", + test_case.description + ); + assert_eq!( + editable, test_case.expected_editable, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-cache/src/digest/tests/test_types.rs b/crates/sqlbuild-cache/src/digest/tests/test_types.rs index adae6506da..4750a48984 100644 --- a/crates/sqlbuild-cache/src/digest/tests/test_types.rs +++ b/crates/sqlbuild-cache/src/digest/tests/test_types.rs @@ -50,3 +50,12 @@ pub(crate) struct CompiledCodeIdentityTestCase { pub(crate) build_identity: &'static str, pub(crate) expected_identity: &'static str, } + +/// Search path folders with their files, and the facts `importlib.metadata` reads from them. +pub(crate) struct InstalledCodeFactsTestCase { + pub(crate) description: &'static str, + /// `(search path index, relative path, contents)`. + pub(crate) files: &'static [(usize, &'static str, &'static str)], + pub(crate) expected_versions: [&'static str; 3], + pub(crate) expected_editable: bool, +} diff --git a/src/sqlbuild/compiler/profiling/types.py b/crates/sqlbuild-core/src/compile_timings/constants.rs similarity index 66% rename from src/sqlbuild/compiler/profiling/types.py rename to crates/sqlbuild-core/src/compile_timings/constants.rs index 59ff6ef9d8..00f84cac18 100644 --- a/src/sqlbuild/compiler/profiling/types.py +++ b/crates/sqlbuild-core/src/compile_timings/constants.rs @@ -1,8 +1,7 @@ -"""Compile profiling type aliases.""" +//! The compile phases and metrics `--json` reports, in presentation order. -from typing import Literal - -type CompileTimingPhase = Literal[ +/// Phases whose elapsed nanoseconds are reported as whole milliseconds. +pub const COMPILE_TIMING_PHASES: [&str; 13] = [ "attachment_ms", "model_analysis_ms", "model_analysis_cpu_ms", @@ -16,13 +15,14 @@ "cache_publication_ms", "physical_write_ms", "stale_traversal_ms", -] +]; -type CompileMetric = Literal[ +/// Counters reported as recorded. +pub const COMPILE_METRICS: [&str; 6] = [ "analysis_batch_cache_hits", "analysis_entry_cache_hits", "analysis_cache_misses", "analysis_cache_bypasses", "sql_test_scan_cache_hits", "sql_test_scan_cache_misses", -] +]; diff --git a/crates/sqlbuild-core/src/compile_timings/main/compile_timing_milliseconds.rs b/crates/sqlbuild-core/src/compile_timings/main/compile_timing_milliseconds.rs new file mode 100644 index 0000000000..e71a01e4cb --- /dev/null +++ b/crates/sqlbuild-core/src/compile_timings/main/compile_timing_milliseconds.rs @@ -0,0 +1,19 @@ +//! Every supported phase in milliseconds, then every metric, in stable presentation order. + +use crate::compile_timings::constants::{COMPILE_METRICS, COMPILE_TIMING_PHASES}; +use crate::compile_timings::models::CompileTimings; + +/// Each phase's whole milliseconds and each metric's value, absent ones as zero. +#[must_use] +pub fn compile_timing_milliseconds(timings: &CompileTimings) -> Vec<(&'static str, u64)> { + let phases = COMPILE_TIMING_PHASES.iter().map(|phase| { + ( + *phase, + timings.elapsed_ns.get(*phase).copied().unwrap_or(0) / 1_000_000, + ) + }); + let metrics = COMPILE_METRICS + .iter() + .map(|metric| (*metric, timings.metrics.get(*metric).copied().unwrap_or(0))); + phases.chain(metrics).collect() +} diff --git a/crates/sqlbuild-core/src/compile_timings/main/mod.rs b/crates/sqlbuild-core/src/compile_timings/main/mod.rs new file mode 100644 index 0000000000..10e1c9c2ee --- /dev/null +++ b/crates/sqlbuild-core/src/compile_timings/main/mod.rs @@ -0,0 +1 @@ +pub mod compile_timing_milliseconds; diff --git a/crates/sqlbuild-core/src/compile_timings/mod.rs b/crates/sqlbuild-core/src/compile_timings/mod.rs new file mode 100644 index 0000000000..e13b4e8dfa --- /dev/null +++ b/crates/sqlbuild-core/src/compile_timings/mod.rs @@ -0,0 +1,8 @@ +//! Detailed timings and cache metrics one compile command accumulates for `--json` output. + +pub mod constants; +pub mod main; +pub mod models; + +#[cfg(test)] +mod tests; diff --git a/crates/sqlbuild-core/src/compile_timings/models.rs b/crates/sqlbuild-core/src/compile_timings/models.rs new file mode 100644 index 0000000000..ba95cdb4c8 --- /dev/null +++ b/crates/sqlbuild-core/src/compile_timings/models.rs @@ -0,0 +1,10 @@ +//! Accumulated compile timings. + +use std::collections::HashMap; + +/// Elapsed nanoseconds per phase and summed values per metric of one compile invocation. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct CompileTimings { + pub elapsed_ns: HashMap, + pub metrics: HashMap, +} diff --git a/crates/sqlbuild-core/src/compile_timings/tests.rs b/crates/sqlbuild-core/src/compile_timings/tests.rs new file mode 100644 index 0000000000..191b814f97 --- /dev/null +++ b/crates/sqlbuild-core/src/compile_timings/tests.rs @@ -0,0 +1,4 @@ +#[path = "tests/test_compile_timing_milliseconds.rs"] +mod compile_timing_milliseconds; +#[path = "tests/test_types.rs"] +mod test_types; diff --git a/crates/sqlbuild-core/src/compile_timings/tests/test_compile_timing_milliseconds.rs b/crates/sqlbuild-core/src/compile_timings/tests/test_compile_timing_milliseconds.rs new file mode 100644 index 0000000000..ea3cfee116 --- /dev/null +++ b/crates/sqlbuild-core/src/compile_timings/tests/test_compile_timing_milliseconds.rs @@ -0,0 +1,37 @@ +use std::collections::HashMap; + +use crate::compile_timings::main::compile_timing_milliseconds::compile_timing_milliseconds; +use crate::compile_timings::models::CompileTimings; +use crate::compile_timings::tests::test_types::MillisecondsTestCase; + +#[test] +fn given_recorded_timings_when_reported_then_lists_every_phase_and_metric_in_order() { + let test_cases = [ + MillisecondsTestCase { + description: "nothing recorded reports zeros", + timings: CompileTimings::default(), + expected_entries: vec![("attachment_ms", 0), ("sql_test_scan_cache_misses", 0)], + }, + MillisecondsTestCase { + description: "nanoseconds floor to milliseconds and unknown phases are dropped", + timings: CompileTimings { + elapsed_ns: HashMap::from([ + ("attachment_ms".to_owned(), 2_999_999), + ("unknown_ms".to_owned(), 5_000_000), + ]), + metrics: HashMap::from([("sql_test_scan_cache_misses".to_owned(), 3)]), + }, + expected_entries: vec![("attachment_ms", 2), ("sql_test_scan_cache_misses", 3)], + }, + ]; + for test_case in test_cases { + let entries = compile_timing_milliseconds(&test_case.timings); + assert_eq!(entries.len(), 19, "{}", test_case.description); + assert_eq!( + vec![entries[0], entries[18]], + test_case.expected_entries, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-core/src/compile_timings/tests/test_types.rs b/crates/sqlbuild-core/src/compile_timings/tests/test_types.rs new file mode 100644 index 0000000000..db51d06fbc --- /dev/null +++ b/crates/sqlbuild-core/src/compile_timings/tests/test_types.rs @@ -0,0 +1,7 @@ +use crate::compile_timings::models::CompileTimings; + +pub(super) struct MillisecondsTestCase { + pub(super) description: &'static str, + pub(super) timings: CompileTimings, + pub(super) expected_entries: Vec<(&'static str, u64)>, +} diff --git a/crates/sqlbuild-core/src/lib.rs b/crates/sqlbuild-core/src/lib.rs index f2911bf659..bbb65f2801 100644 --- a/crates/sqlbuild-core/src/lib.rs +++ b/crates/sqlbuild-core/src/lib.rs @@ -9,6 +9,7 @@ pub mod build_progress; pub mod check_command; pub mod cli_text; pub mod clone_text; +pub mod compile_timings; pub mod constants; pub mod contract_command; pub mod cost_allocation; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/model_analyses.rs b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/model_analyses.rs index 911eff0bb4..fdb93624a8 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/model_analyses.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/analysis_session/model_analyses.rs @@ -11,13 +11,13 @@ use crate::bindings::_helpers::analysis_session::profile_inputs::{ use crate::bindings::_helpers::analysis_session::stage_functions::family_rows; use crate::bindings::_helpers::attachments::cursor_intrinsics::analysis_cursor_sql; use crate::bindings::_helpers::boundary::keyword_request::Request; +use crate::bindings::_helpers::cache::native_store::current_code_identity; use crate::bindings::_helpers::compile_driver::analysis_cache::record_analysis_cache_metrics; const NATIVE_MODULE: &str = "sqlbuild._native"; const ANALYSIS_MODELS_MODULE: &str = "sqlbuild.compiler.analysis_session.models"; const PLANNER_TYPES_MODULE: &str = "sqlbuild.compiler.planner.types"; const LINEAGE_TYPES_MODULE: &str = "sqlbuild.compiler.lineage.types"; -const CODE_IDENTITY_MODULE: &str = "sqlbuild.compiler.frontier.main.compiled_code_identity"; const LOG_DEBUG_MODULE: &str = "sqlbuild.diagnostics.main.log_debug_event"; const LOGGING_MODULE: &str = "logging"; const DEBUG_LOGGER: &str = "sqlbuild.compile"; @@ -119,10 +119,7 @@ fn analysis_store<'py>(request: &Bound<'py, PyAny>) -> PyResult = py - .import(CODE_IDENTITY_MODULE)? - .getattr("compiled_code_identity")? - .call0()?; + let identity: Bound<'py, PyAny> = PyString::new(py, ¤t_code_identity(py)?).into_any(); let parts: Bound<'py, PyList> = PyList::new( py, [ @@ -259,8 +256,12 @@ fn record_cache(python: &[Bound<'_, PyAny>], session: &Bound<'_, PyAny>) -> PyRe } /// Record the analysis cache entries a native session read, missed and bypassed. -fn record_metrics(py: Python<'_>, (hits, misses, bypasses): (usize, usize, usize)) -> PyResult<()> { - record_analysis_cache_metrics(py, (0, hits, misses, bypasses)) +fn record_metrics( + _py: Python<'_>, + (hits, misses, bypasses): (usize, usize, usize), +) -> PyResult<()> { + record_analysis_cache_metrics((0, hits, misses, bypasses)); + Ok(()) } /// Each model's native request row, in analysis order. diff --git a/crates/sqlbuild-python/src/bindings/_helpers/cache/native_store.rs b/crates/sqlbuild-python/src/bindings/_helpers/cache/native_store.rs index 090b854fed..792448b1b7 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/cache/native_store.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/cache/native_store.rs @@ -1,22 +1,28 @@ //! Digests, project fingerprints and store files of the shared native store, with the GIL released. use pyo3::exceptions::{PyOSError, PyRuntimeError}; -use pyo3::prelude::{Bound, PyModule, PyModuleMethods, PyResult, Python}; +use pyo3::prelude::{Bound, PyAnyMethods, PyErr, PyModule, PyModuleMethods, PyResult, Python}; +use pyo3::sync::PyOnceLock; use pyo3::{pyfunction, wrap_pyfunction}; use sqlbuild_cache::digest::main::compiled_code_identity::compiled_code_identity as compiled; use sqlbuild_cache::digest::main::content_digest::content_digest as digest; use sqlbuild_cache::digest::main::digest_files::digest_files as digest_each; use sqlbuild_cache::digest::main::fingerprint_project_files::fingerprint_project_files as fingerprint; use sqlbuild_cache::digest::main::hex_digest::hex_digest; +use sqlbuild_cache::digest::main::installed_code_facts::installed_code_facts; use sqlbuild_cache::digest::main::installed_code_identity::installed_code_identity as installed; use sqlbuild_cache::digest::types::ContentDigest; use sqlbuild_cache::store::main::open_native_store::open_native_store; use sqlbuild_cache::store::models::NativeStore; +use sqlbuild_rules::constants::NATIVE_BUILD_IDENTITY; use std::collections::HashSet; use std::path::{Path, PathBuf}; use crate::bindings::types::CompilerDetach; +const NATIVE_MODULE: &str = "sqlbuild._native"; +const PACKAGE_MODULE: &str = "sqlbuild"; + /// Hexadecimal digest of length-prefixed UTF-8 `parts`, the digest every native cache key uses. #[pyfunction] pub(crate) fn content_digest(parts: Vec) -> String { @@ -82,25 +88,69 @@ pub(crate) fn save_store( .map_err(|error| PyOSError::new_err(error.to_string())) } -/// Digest of the installed code's facts and, for an editable install, its source files. +/// The digest that changes with every released, locally edited, or rebuilt SQLBuild version. #[pyfunction] -pub(crate) fn installed_code_identity( - py: Python<'_>, - facts: Vec, - package_root: Option, -) -> PyResult { - py.compiler_detach(|| Ok(installed(&facts, package_root.as_deref()))) - .map_err(PyRuntimeError::new_err) +pub(crate) fn compiled_code_identity(py: Python<'_>) -> PyResult { + static IDENTITY: PyOnceLock = PyOnceLock::new(); + IDENTITY + .get_or_try_init(py, || { + Ok::(compiled(&installed_identity(py)?, NATIVE_BUILD_IDENTITY)) + }) + .cloned() } -/// Digest of the installed code's identity and the native build's identity. -#[pyfunction] -pub(crate) fn compiled_code_identity(installed: &str, build_identity: &str) -> String { - compiled(installed, build_identity) +/// `compiled_code_identity`, looked up on the extension module so tests can replace it. +pub(crate) fn current_code_identity(py: Python<'_>) -> PyResult { + py.import(NATIVE_MODULE)? + .getattr("compiled_code_identity")? + .call0()? + .extract() +} + +/// Installed distribution versions, the interpreter, and an editable install's source files. +fn installed_identity(py: Python<'_>) -> PyResult { + let search_paths: Vec = py + .import("sys")? + .getattr("path")? + .try_iter()? + .map(|entry| entry?.extract::()) + .collect::>()?; + let implementation: String = py + .import("platform")? + .call_method0("python_implementation")? + .extract()?; + let (major, minor): (u32, u32) = ( + py.import("sys")? + .getattr("version_info")? + .getattr("major")? + .extract()?, + py.import("sys")? + .getattr("version_info")? + .getattr("minor")? + .extract()?, + ); + let package_file: PathBuf = py.import(PACKAGE_MODULE)?.getattr("__file__")?.extract()?; + let (mut facts, editable) = py + .compiler_detach(|| Ok(installed_code_facts(&search_paths))) + .map_err(PyRuntimeError::new_err)?; + facts.push(implementation); + facts.push(format!("{major}.{minor}")); + let package_root: Option = if editable { + Some( + package_file + .parent() + .unwrap_or(&package_file) + .canonicalize() + .map_err(|error| PyOSError::new_err(error.to_string()))?, + ) + } else { + None + }; + py.compiler_detach(|| Ok(installed(&facts, package_root.as_deref()))) + .map_err(PyRuntimeError::new_err) } pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add_function(wrap_pyfunction!(installed_code_identity, module)?)?; module.add_function(wrap_pyfunction!(compiled_code_identity, module)?)?; module.add_function(wrap_pyfunction!(content_digest, module)?)?; module.add_function(wrap_pyfunction!(fingerprint_project_files, module)?)?; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/cache/sql_test_scan_store.rs b/crates/sqlbuild-python/src/bindings/_helpers/cache/sql_test_scan_store.rs index a315e0f06a..a5d7a6f9b5 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/cache/sql_test_scan_store.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/cache/sql_test_scan_store.rs @@ -2,7 +2,6 @@ use pyo3::exceptions::PyValueError; use pyo3::prelude::{Bound, Py, PyAnyMethods, PyModule, PyModuleMethods, PyResult, Python}; -use pyo3::types::{PyDict, PyDictMethods}; use pyo3::{pyclass, pymethods}; use sqlbuild_cache::digest::main::content_digest::content_digest; use sqlbuild_cache::digest::main::hex_digest::hex_digest; @@ -12,11 +11,11 @@ use std::path::PathBuf; use std::sync::Mutex; use std::sync::atomic::{AtomicUsize, Ordering}; +use crate::bindings::_helpers::cache::native_store::current_code_identity; use crate::bindings::_helpers::cache::native_store::{open_store, save_store}; +use crate::bindings::_helpers::compile_driver::compile_timings::record_compile_metric; use crate::bindings::_helpers::compile_driver::python_calls::imported; -const CODE_IDENTITY_MODULE: &str = "sqlbuild.compiler.frontier.main.compiled_code_identity"; -const METRIC_MODULE: &str = "sqlbuild.compiler.profiling.main._metric"; const SCAN_CACHE_LOGGER: &str = "sqlbuild.compile"; /// The native store kind holding SQL-test scan results. @@ -140,9 +139,7 @@ impl SqlTestScanCache { let Some(directory) = cache_dir.as_ref().filter(|_| store.is_none()) else { return Ok(store.as_ref().map(|store| store.clone_ref(py))); }; - let identity: String = imported(py, CODE_IDENTITY_MODULE, "compiled_code_identity")? - .call0()? - .extract()?; + let identity: String = current_code_identity(py)?; let (opened, error) = SqlTestScanStore::open_in(py, directory.clone(), &identity); *store = opened.map(|opened| Py::new(py, opened)).transpose()?; if let Some(error) = error { @@ -162,8 +159,14 @@ impl SqlTestScanCache { .as_ref() .map(|store| store.clone_ref(py)); let (hits, misses) = store.as_ref().map_or((0, 0), |store| store.get().counts()); - metric(py, "sql_test_scan_cache_hits", hits)?; - metric(py, "sql_test_scan_cache_misses", misses)?; + record_compile_metric( + "sql_test_scan_cache_hits", + i64::try_from(hits).unwrap_or(i64::MAX), + ); + record_compile_metric( + "sql_test_scan_cache_misses", + i64::try_from(misses).unwrap_or(i64::MAX), + ); let Some(store) = store else { return Ok(()); }; @@ -174,14 +177,6 @@ impl SqlTestScanCache { } } -fn metric(py: Python<'_>, name: &str, value: usize) -> PyResult<()> { - let keywords: Bound<'_, PyDict> = PyDict::new(py); - keywords.set_item("metric", name)?; - keywords.set_item("value", value)?; - imported(py, METRIC_MODULE, "record_compile_metric")?.call((), Some(&keywords))?; - Ok(()) -} - fn debug(py: Python<'_>, message: &str, error: &str) -> PyResult<()> { imported(py, "logging", "getLogger")? .call1((SCAN_CACHE_LOGGER,))? diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compile_command/artifact_batch.rs b/crates/sqlbuild-python/src/bindings/_helpers/compile_command/artifact_batch.rs index a9d7ad48f4..9724ba59eb 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compile_command/artifact_batch.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compile_command/artifact_batch.rs @@ -7,23 +7,9 @@ use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyResult, Python}; use crate::bindings::_helpers::compile_command::reuse_recording::{ ArtifactRecorder, artifact_written, }; -use crate::bindings::_helpers::compile_driver::python_calls::{imported, within}; +use crate::bindings::_helpers::compile_driver::compile_timings::record_compile_timing; use crate::bindings::_helpers::compile_outputs::artifacts::write_compiled_artifacts; -const PROFILING_MODULE: &str = "sqlbuild.compiler.profiling.main._record"; - -/// Run `body` inside `record_compile_timing(phase)`. -pub(crate) fn timed<'py, T>( - py: Python<'py>, - phase: &str, - body: impl FnOnce() -> PyResult, -) -> PyResult { - within( - &imported(py, PROFILING_MODULE, "record_compile_timing")?.call1((phase,))?, - |_| body(), - ) -} - /// The artifacts of one write group, written together by `flush`. pub(crate) struct ArtifactBatch<'a> { check_existing: bool, @@ -74,7 +60,7 @@ impl<'a> ArtifactBatch<'a> { return Ok(()); } let files: Vec<(PathBuf, Vec)> = std::mem::take(&mut self.files); - timed(py, "physical_write_ms", || { + record_compile_timing("physical_write_ms", || { write_compiled_artifacts(py, files, self.check_existing).map(|_| ()) }) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compile_command/target_writer.rs b/crates/sqlbuild-python/src/bindings/_helpers/compile_command/target_writer.rs index cd267b1829..a71989a2db 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compile_command/target_writer.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compile_command/target_writer.rs @@ -14,7 +14,7 @@ use sqlbuild_cache::artifacts::main::reusable_sql_test_record::reusable_sql_test use sqlbuild_cache::artifacts::models::SqlTestArtifactRecord; use crate::bindings::_helpers::analysis::sql_test_comparisons::single_comparison_sql; -use crate::bindings::_helpers::compile_command::artifact_batch::{ArtifactBatch, timed}; +use crate::bindings::_helpers::compile_command::artifact_batch::ArtifactBatch; use crate::bindings::_helpers::compile_command::reuse_recording::{ ArtifactRecorder, CompileArtifactWrites, artifact_kept, artifact_moved, artifact_tree_moved, }; @@ -22,6 +22,7 @@ use crate::bindings::_helpers::compile_command::test_artifact_cache::{ IdentityContext, artifact_identity, artifact_record_key, identity_context, read_cache, write_cache, }; +use crate::bindings::_helpers::compile_driver::compile_timings::record_compile_timing; use crate::bindings::_helpers::compile_driver::python_calls::{ KeywordArguments, call_with, imported, keywords, within, }; @@ -154,7 +155,7 @@ fn remove_stale(target_dir: &Bound<'_, PyAny>, managed: &[Bound<'_, PyAny>]) -> .iter() .map(|path| path.extract()) .collect::>()?; - timed(py, "stale_traversal_ms", || { + record_compile_timing("stale_traversal_ms", || { remove_stale_artifacts(py, compiled, managed).map(|_| ()) }) } @@ -270,7 +271,7 @@ fn write_planned_target<'py>( let path: Bound<'py, PyAny> = compiled .div(TESTS_DIR)? .div(planned_sql_test_path(&entry)?)?; - let comparison_sql: String = timed(py, "comparison_render_ms", || { + let comparison_sql: String = record_compile_timing("comparison_render_ms", || { single_comparison_sql(&entry, adapter) })?; batch.queue_sql(&path, &comparison_sql)?; @@ -547,7 +548,7 @@ fn write_static_tests<'py>( } } } - timed(py, "cache_publication_ms", || { + record_compile_timing("cache_publication_ms", || { write_cache(&project.getattr("compile_cache_dir")?, &records) })?; Ok((managed, diagnostics)) @@ -651,7 +652,7 @@ pub(crate) fn publish_staged_target<'py>( .call1((expected_files,))? .extract()?; let (staged, compiled): (PathBuf, PathBuf) = (staged_dir.extract()?, compiled_dir.extract()?); - let (kind, published) = timed(py, "physical_write_ms", || { + let (kind, published) = record_compile_timing("physical_write_ms", || { publish_staged_artifacts(py, staged, compiled, expected) })?; if kind == PUBLISHED_STAGING_CHANGED { diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/analysis_cache.rs b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/analysis_cache.rs index a597342342..4028306fef 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/analysis_cache.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/analysis_cache.rs @@ -8,6 +8,7 @@ use pyo3::types::{PyBool, PyDict, PyDictMethods, PyList, PyListMethods, PySlice, use pyo3::{pyfunction, wrap_pyfunction}; use sqlbuild_cache::digest::main::sha256_hex::sha256_hex; +use crate::bindings::_helpers::compile_driver::compile_timings::record_compile_metric; use crate::bindings::_helpers::compile_driver::python_calls::{ KeywordArguments, call_with, imported, }; @@ -21,7 +22,6 @@ const UNKNOWN_VERSION: &str = "unknown"; const COMPILE_MODELS_MODULE: &str = "sqlbuild.compiler.compile.models"; /// Sorted; the dialects whose expression types compile checks. const TYPE_CHECKED_DIALECTS: [&str; 4] = ["bigquery", "duckdb", "postgresql", "snowflake"]; -const METRIC_MODULE: &str = "sqlbuild.compiler.profiling.main._metric"; const METADATA_MODULE: &str = "importlib.metadata"; const CONTEXT_REQUIRED: [&str; 3] = ["root", "inference_profile", "allow_compact_analysis"]; const CONTEXT_OPTIONAL: [&str; 2] = ["rich_type_inference", "signature_namespace"]; @@ -121,22 +121,16 @@ fn build_analysis_cache_context<'py>( /// Record semantic analysis cache outcomes for one invocation. pub(crate) fn record_analysis_cache_metrics( - py: Python<'_>, (batch_hits, entry_hits, misses, bypasses): (usize, usize, usize, usize), -) -> PyResult<()> { - let record: Bound<'_, PyAny> = imported(py, METRIC_MODULE, "record_compile_metric")?; +) { for (metric, value) in [ ("analysis_batch_cache_hits", batch_hits), ("analysis_entry_cache_hits", entry_hits), ("analysis_cache_misses", misses), ("analysis_cache_bypasses", bypasses), ] { - let keywords: Bound<'_, PyDict> = PyDict::new(py); - keywords.set_item("metric", metric)?; - keywords.set_item("value", value)?; - record.call((), Some(&keywords))?; + record_compile_metric(metric, i64::try_from(value).unwrap_or(i64::MAX)); } - Ok(()) } /// Whether `error` is the `TypeError` or `ValueError` an unstable payload raises. diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/assembly_analysis.rs b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/assembly_analysis.rs index 8114cefb8b..6f79a80240 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/assembly_analysis.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/assembly_analysis.rs @@ -8,15 +8,14 @@ use sqlbuild_analysis::assembly::analysis_session::main::native_dialect::native_ use crate::bindings::_helpers::analysis::binding_catalog::{BindingCatalog, ignores_quoted_case}; use crate::bindings::_helpers::compile_driver::analysis_cache::analysis_cache_context; use crate::bindings::_helpers::compile_driver::assembly::AssemblyRequest; -use crate::bindings::_helpers::compile_driver::python_calls::{ - call_with, imported, keywords, within, +use crate::bindings::_helpers::compile_driver::compile_timings::{ + record_compile_cpu_timing, record_compile_timing, }; +use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported, keywords}; const NATIVE_MODULE: &str = "sqlbuild._native"; const ADAPTER_MODELS_MODULE: &str = "sqlbuild.adapter.contract.models"; pub(crate) const ANALYSIS_MODELS_MODULE: &str = "sqlbuild.compiler.analysis_session.models"; -const PROFILING_MODULE: &str = "sqlbuild.compiler.profiling.main._record"; -const PROFILING_CPU_MODULE: &str = "sqlbuild.compiler.profiling.main._record_cpu"; const LINEAGE_TYPES_MODULE: &str = "sqlbuild.compiler.lineage.types"; const PLANNER_TYPES_MODULE: &str = "sqlbuild.compiler.planner.types"; const GENERIC_DIALECT: &str = "generic"; @@ -347,13 +346,8 @@ fn model_analyses<'py>( } else { py.None().into_bound(py) }; - let timing: Bound<'py, PyAny> = - imported(py, PROFILING_MODULE, "record_compile_timing")?.call1(("model_analysis_ms",))?; - within(&timing, |_| { - let cpu_timing: Bound<'py, PyAny> = - imported(py, PROFILING_CPU_MODULE, "record_compile_cpu_timing")? - .call1(("model_analysis_cpu_ms",))?; - within(&cpu_timing, |_| { + record_compile_timing("model_analysis_ms", || { + record_compile_cpu_timing(py, "model_analysis_cpu_ms", || { let (known_functions, known_types): (Bound<'py, PyAny>, Bound<'py, PyAny>) = declared_names(inputs, &context.column_types)?; let mut model_inputs: Vec> = Vec::new(); diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/command.rs b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/command.rs index c2cdee3133..445f559674 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/command.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/command.rs @@ -3,6 +3,7 @@ use pyo3::prelude::{Bound, Py, PyAny, PyAnyMethods, PyModule, PyModuleMethods, PyResult, Python}; use pyo3::types::{PyBool, PyDict, PyDictMethods, PyMapping, PyString, PyTuple}; use pyo3::{IntoPyObject, pyfunction, wrap_pyfunction}; +use std::rc::Rc; use crate::bindings::_helpers::build_progress::transient_status::status_reporter; use crate::bindings::_helpers::compile_command::models::CompileReuseAttempt; @@ -13,6 +14,9 @@ use crate::bindings::_helpers::compile_command::reuse_recording::{ CompileArtifactWrites, RecordedCompileOutput, compile_reuse_timings, new_artifact_recorder, }; use crate::bindings::_helpers::compile_driver::analysis::analyze; +use crate::bindings::_helpers::compile_driver::compile_timings::{ + CompileTimingsHandle, collect_compile_timings, timing_milliseconds, +}; use crate::bindings::_helpers::compile_driver::phases::{ elapsed_ms, manifest_payload, monotonic, resolved_dag_path, sql_test_planning_diagnostics, target_dir, write_artifacts, write_dag_artifact, @@ -27,7 +31,6 @@ const BYPASS_OUTCOME: &str = "bypass"; const NATIVE_MODULE: &str = "sqlbuild._native"; const COMPILE_MODELS_MODULE: &str = "sqlbuild.compiler.compile.models"; const COMPILE_TYPES_MODULE: &str = "sqlbuild.compiler.compile.types"; -const PROFILING_MODULE: &str = "sqlbuild.compiler.profiling.main._collect"; /// What one full compile reports for reuse recording. struct CommandOutcome<'py> { @@ -39,7 +42,7 @@ struct CommandOutcome<'py> { /// The per-invocation recorders a compile runs inside. struct Recorders<'py> { - detailed_timings: Bound<'py, PyAny>, + detailed_timings: CompileTimingsHandle, input_reads: Bound<'py, PyAny>, artifact_writes: Py, prepared_artifacts: Bound<'py, PyAny>, @@ -146,9 +149,8 @@ fn recorded_compile<'py>( .getattr("profile_flags")? .getattr("skip_write")? .is_truthy()?; - let compiled: PyResult<(CommandOutcome<'py>, Bound<'py, PyAny>, Bound<'py, PyAny>)> = within( - &imported(py, PROFILING_MODULE, "collect_compile_timings")?.call0()?, - |detailed_timings| { + let compiled: PyResult<(CommandOutcome<'py>, Bound<'py, PyAny>, Bound<'py, PyAny>)> = + collect_compile_timings(|detailed_timings| { within( &imported(py, NATIVE_MODULE, "record_compile_input_reads")?.call0()?, |input_reads| { @@ -168,7 +170,7 @@ fn recorded_compile<'py>( )?; within(&manager, |prepared_artifacts| { let recorders: Recorders<'py> = Recorders { - detailed_timings: detailed_timings.clone(), + detailed_timings: Rc::clone(detailed_timings), input_reads: input_reads.clone(), artifact_writes: artifact_writes.clone_ref(py), prepared_artifacts: prepared_artifacts.clone(), @@ -188,8 +190,7 @@ fn recorded_compile<'py>( }) }, ) - }, - ); + }); close_status(status)?; let (outcome, input_reads, artifact_writes) = compiled?; let stored: PyResult> = call_with( @@ -372,7 +373,7 @@ fn rules_failed(analysis: &Bound<'_, PyAny>) -> PyResult { /// The compile's phase timings, in the order `--json` reports them. fn timings_ms<'py>( (analysis, write_ms): (&Bound<'py, PyAny>, i64), - (detailed_timings, reuse_timings): (&Bound<'py, PyAny>, &Bound<'py, PyAny>), + (detailed_timings, reuse_timings): (&CompileTimingsHandle, &Bound<'py, PyAny>), total_start: f64, ) -> PyResult> { let py: Python<'py> = analysis.py(); @@ -393,11 +394,9 @@ fn timings_ms<'py>( timings.set_item(key, analysis.getattr(attribute)?)?; } timings.set_item("write_ms", write_ms)?; - timings.update( - detailed_timings - .call_method0("as_milliseconds")? - .downcast::()?, - )?; + for (key, value) in timing_milliseconds(detailed_timings) { + timings.set_item(key, value)?; + } timings.update(reuse_timings.downcast::()?)?; timings.set_item("total_ms", elapsed_ms(py, total_start)?)?; Ok(timings) diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/compile_timings.rs b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/compile_timings.rs new file mode 100644 index 0000000000..869a6289d7 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/compile_timings.rs @@ -0,0 +1,98 @@ +//! The active compile timing collector of the current thread, as a context variable held it. + +use std::cell::RefCell; +use std::rc::Rc; +use std::time::Instant; + +use pyo3::prelude::{PyAnyMethods, PyResult, Python}; +use sqlbuild_core::compile_timings::main::compile_timing_milliseconds::compile_timing_milliseconds; +use sqlbuild_core::compile_timings::models::CompileTimings; + +use crate::bindings::_helpers::compile_driver::python_calls::imported; + +/// One compile invocation's collector, shared by every recorder on its thread. +pub(crate) type CompileTimingsHandle = Rc>; + +thread_local! { + static ACTIVE: RefCell> = const { RefCell::new(None) }; +} + +/// Run `body` with a fresh collector active on this thread, restoring the previous one after. +pub(crate) fn collect_compile_timings( + body: impl FnOnce(&CompileTimingsHandle) -> PyResult, +) -> PyResult { + let collector: CompileTimingsHandle = Rc::default(); + let previous: Option = + ACTIVE.with(|active| active.replace(Some(Rc::clone(&collector)))); + let result: PyResult = body(&collector); + ACTIVE.with(|active| active.replace(previous)); + result +} + +/// Every reported phase in milliseconds, then every metric. +pub(crate) fn timing_milliseconds(collector: &CompileTimingsHandle) -> Vec<(&'static str, u64)> { + compile_timing_milliseconds(&collector.borrow()) +} + +/// Add elapsed nanoseconds to a phase when a collector is active. +pub(crate) fn add_compile_timing(phase: &str, elapsed_ns: u64) { + if let Some(collector) = active() { + *collector + .borrow_mut() + .elapsed_ns + .entry(phase.to_owned()) + .or_default() += elapsed_ns; + } +} + +/// Add a non-negative metric value when a collector is active. +pub(crate) fn record_compile_metric(metric: &str, value: i64) { + if let Some(collector) = active() { + *collector + .borrow_mut() + .metrics + .entry(metric.to_owned()) + .or_default() += u64::try_from(value).unwrap_or(0); + } +} + +/// Record the wall time `body` spans under `phase` when a collector is active. +pub(crate) fn record_compile_timing( + phase: &str, + body: impl FnOnce() -> PyResult, +) -> PyResult { + if active().is_none() { + return body(); + } + let started: Instant = Instant::now(); + let result: PyResult = body(); + add_compile_timing( + phase, + u64::try_from(started.elapsed().as_nanos()).unwrap_or(u64::MAX), + ); + result +} + +/// Record the process CPU time `body` spans, including its worker threads, under `phase`. +pub(crate) fn record_compile_cpu_timing( + py: Python<'_>, + phase: &str, + body: impl FnOnce() -> PyResult, +) -> PyResult { + if active().is_none() { + return body(); + } + let started: u64 = process_time_ns(py)?; + let result: PyResult = body(); + let ended: u64 = process_time_ns(py)?; + add_compile_timing(phase, ended.saturating_sub(started)); + result +} + +fn process_time_ns(py: Python<'_>) -> PyResult { + imported(py, "time", "process_time_ns")?.call0()?.extract() +} + +fn active() -> Option { + ACTIVE.with(|active| active.borrow().clone()) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/mod.rs index dabff08ac1..cd7dafa09c 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/mod.rs @@ -8,6 +8,7 @@ pub(crate) mod assembly_models; pub(crate) mod assembly_project; pub(crate) mod assembly_sql_tests; pub(crate) mod command; +pub(crate) mod compile_timings; pub(crate) mod config_validation; pub(crate) mod deferral; pub(crate) mod diagnostic_collector; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/project.rs b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/project.rs index a180c5e8ab..e4f8b44114 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/project.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/project.rs @@ -4,10 +4,11 @@ use pyo3::prelude::{Bound, Py, PyAny, PyAnyMethods, PyModule, PyModuleMethods, P use pyo3::types::{PyBool, PyCFunction, PyDict, PyTuple}; use pyo3::{pyfunction, wrap_pyfunction}; +use crate::bindings::_helpers::compile_driver::compile_timings::record_compile_timing; use crate::bindings::_helpers::compile_driver::phases::monotonic; use crate::bindings::_helpers::compile_driver::pipeline_entries::project_graph; use crate::bindings::_helpers::compile_driver::python_calls::{ - KeywordArguments, call_with, imported, keywords, within, + KeywordArguments, call_with, imported, keywords, }; use crate::bindings::_helpers::compile_driver::target_defaults::{ TargetDefaults, apply_target_defaults, @@ -28,7 +29,6 @@ const DISCOVERY_CONSTANTS_MODULE: &str = "sqlbuild.compiler.discovery.constants" const LOCAL_CONFIG_FILENAME: &str = "sqlbuild_local.toml"; const LINEAGE_TYPES_MODULE: &str = "sqlbuild.compiler.lineage.types"; const ANALYSIS_SELECTION_MODULE: &str = "sqlbuild._native"; -const PROFILING_MODULE: &str = "sqlbuild.compiler.profiling.main._record"; const ADAPTER_NAME_MODULE: &str = "sqlbuild._native"; const COLLECTION_RENDERING_MODULE: &str = "sqlbuild._native"; /// Optional keywords of the project graph builders. @@ -152,10 +152,9 @@ fn python_compiled_project<'py>( managed_loader_target_isolation(discovered, adapter, request.resolved_connection.bind(py))?; let selection: &Bound<'py, PyAny> = request.analysis_selection.bind(py); let no_cache: bool = !selection.is_none() && selection.getattr("no_cache")?.is_truthy()?; - let timing: Bound<'py, PyAny> = - imported(py, PROFILING_MODULE, "record_compile_timing")?.call1(("attachment_ms",))?; - let compile_inputs: Bound<'py, PyAny> = - within(&timing, |_| frontier_compile_inputs(py, request, no_cache))?; + let compile_inputs: Bound<'py, PyAny> = record_compile_timing("attachment_ms", || { + frontier_compile_inputs(py, request, no_cache) + })?; if !selection.is_none() { let on_inputs_ready: Bound<'py, PyAny> = selection.getattr("on_inputs_ready")?; if !on_inputs_ready.is_none() { diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/attachment.rs b/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/attachment.rs index 4d937e7809..c2d41b0a42 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/attachment.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compiled_project/attachment.rs @@ -5,8 +5,11 @@ use pyo3::types::{PyCFunction, PyDict, PyDictMethods, PyTuple}; use pyo3::{IntoPyObjectExt, pyfunction, wrap_pyfunction}; use std::collections::HashMap; +use sqlbuild_attachments::audits::constants::BUILTIN_AUDITS; + +use crate::bindings::_helpers::compile_driver::compile_timings::record_compile_timing; use crate::bindings::_helpers::compile_driver::python_calls::{ - KeywordArguments, call_with, imported, within, + KeywordArguments, call_with, imported, }; use crate::bindings::_helpers::compiled_project::audit_loop::{AuditBlockRow, AuditFileRow}; use crate::bindings::_helpers::compiled_project::function_loop::FunctionRow; @@ -28,8 +31,7 @@ const COMPILE_TYPES_MODULE: &str = "sqlbuild.compiler.compile.types"; const SCOPE_TYPES_MODULE: &str = "sqlbuild.compiler.scopes.types"; const SCOPE_MODELS_MODULE: &str = "sqlbuild.compiler.scopes.models"; const AUDITING_TYPES_MODULE: &str = "sqlbuild.compiler.auditing.types"; -const BUILTIN_AUDITS_MODULE: &str = "sqlbuild.compiler.auditing.main._builtins"; -const PROFILING_MODULE: &str = "sqlbuild.compiler.profiling.main._record"; +const DISCOVERY_MODELS_MODULE: &str = "sqlbuild.compiler.discovery.models"; const CONTRACT_MODELS_MODULE: &str = "sqlbuild.spec.contracts.models"; const SQL_VALUE_ERRORS_MODULE: &str = "sqlbuild.sql_values.exceptions"; const YAML_FILE_SUFFIXES: [&str; 2] = [".yml", ".yaml"]; @@ -406,6 +408,39 @@ fn audit_file_row(file: Bound<'_, PyAny>) -> PyResult { )) } +/// The discovered file of one built-in generic audit, as project discovery would describe it. +fn builtin_audit_file<'py>( + py: Python<'py>, + (name, sql_body): (&str, &str), +) -> PyResult> { + let path_class: Bound<'py, PyAny> = imported(py, PATHLIB_MODULE, "Path")?; + let relative_path: Bound<'py, PyAny> = + path_class.call1((format!("audits/generic/{name}.sql"),))?; + let block: Bound<'py, PyAny> = call_with( + &imported(py, DISCOVERY_MODELS_MODULE, "DiscoveredAuditBlock")?, + &[ + ("audit_index", 0.into_bound_py_any(py)?), + ("header_values", PyDict::new(py).into_any()), + ("sql_body", sql_body.into_bound_py_any(py)?), + ], + )?; + call_with( + &imported(py, DISCOVERY_MODELS_MODULE, "DiscoveredAuditFile")?, + &[ + ( + "file_path", + path_class.call1(("",))?.div(&relative_path)?, + ), + ("relative_path", relative_path), + ( + "contents", + format!("AUDIT ();\n\n{sql_body}\n").into_bound_py_any(py)?, + ), + ("blocks", PyTuple::new(py, [block])?.into_any()), + ], + ) +} + fn audit_settings<'py>( py: Python<'py>, effective_settings: &Bound<'py, PyAny>, @@ -424,12 +459,10 @@ fn audit_settings<'py>( "default_run_scope", effective_settings.getattr("default_audit_run_scope")?, )?; - let builtins: Vec = - imported(py, BUILTIN_AUDITS_MODULE, "builtin_audit_definitions")? - .call0()? - .try_iter()? - .map(|file| audit_file_row(file?)) - .collect::>()?; + let builtins: Vec = BUILTIN_AUDITS + .iter() + .map(|definition| audit_file_row(builtin_audit_file(py, *definition)?)) + .collect::>()?; settings.set_item("builtins", builtins)?; settings.set_item( "declaration_identity", @@ -584,9 +617,7 @@ pub(crate) fn attach_checks<'py>( ) = arguments.value("settings")?.extract()?; let scan_cache: &Bound<'py, PyAny> = &arguments.value("scan_cache")?; let source_inputs: &Bound<'py, PyAny> = &arguments.value("source_inputs")?; - let timing: Bound<'py, PyAny> = - imported(py, PROFILING_MODULE, "record_compile_timing")?.call1((TEST_TIMING_KEY,))?; - let (tests, scenarios) = within(&timing, |_| { + let (tests, scenarios) = record_compile_timing(TEST_TIMING_KEY, || { let attached = NativeCompiledProject::attach_sql_tests( native_project, items(discovered_inputs, "test_files")? diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/rows.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/rows.rs index 30d5cc9aef..bc993b60e1 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/rows.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/rows.rs @@ -27,6 +27,7 @@ use sqlbuild_rules::rule_cache::main::rules_bulk_cache_path::rules_bulk_cache_pa use crate::bindings::_helpers::boundary::keyword_request::Request; use crate::bindings::_helpers::boundary::panics::{compiler_guard, value_error}; +use crate::bindings::_helpers::cache::native_store::current_code_identity; use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported}; use crate::bindings::_helpers::compiled_project::project::NativeCompiledProject; use crate::bindings::_helpers::compiled_project::records::declaration_rows; @@ -39,7 +40,6 @@ use crate::bindings::types::CompilerDetach; const NATIVE_MODULE: &str = "sqlbuild._native"; const MODELS_MODULE: &str = "sqlbuild.rule_engine.models"; -const CODE_IDENTITY_MODULE: &str = "sqlbuild.compiler.frontier.main.compiled_code_identity"; const DATACLASSES_MODULE: &str = "dataclasses"; const NATIVE_RULES_CACHE_FILE: &str = "native.json"; const NATIVE_RULES_MEMO_FILE: &str = "native-response.json"; @@ -130,9 +130,7 @@ fn built_request(options: &Request<'_, '_>) -> PyResult { let memo: Option<(String, String)> = if enabled { let unresolved: PathBuf = project_dir.extract()?; Some(( - imported(py, CODE_IDENTITY_MODULE, "compiled_code_identity")? - .call0()? - .extract()?, + current_code_identity(py)?, path_text( py, rules_bulk_cache_path(&unresolved, NATIVE_RULES_MEMO_FILE), diff --git a/crates/sqlbuild-python/src/bindings/_helpers/sql_test_glue/plan_driver.rs b/crates/sqlbuild-python/src/bindings/_helpers/sql_test_glue/plan_driver.rs index da46442b4a..1754ec8725 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/sql_test_glue/plan_driver.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/sql_test_glue/plan_driver.rs @@ -1,5 +1,6 @@ //! Host SQL-test request preparation, native planning, and Python result conversion. +use crate::bindings::_helpers::compile_driver::compile_timings::add_compile_timing; use crate::bindings::_helpers::compile_driver::python_calls::{ KeywordArguments, call_with, imported, }; @@ -138,34 +139,18 @@ fn planned<'py>(args: &KeywordArguments<'py>, artifact: bool) -> PyResult> = Vec::new(); for row in native.get_item(0)?.try_iter()? { diff --git a/fensu.toml b/fensu.toml index 15a559828f..8476ce1a4b 100644 --- a/fensu.toml +++ b/fensu.toml @@ -43,6 +43,11 @@ rule = "FFR309" path = "src/sqlbuild/spec/contracts/__init__.py" reason = "Native bindings own configuration resolution; this package only declares the configuration models, constants and errors callers import." +[[targets.python.rule_exceptions]] +rule = "FFR309" +path = "src/sqlbuild/compiler/frontier/__init__.py" +reason = "Native bindings own the compile frontier and the installed code identity; this package only declares the frontier constants, types and errors they read and raise." + [[targets.python.rule_exceptions]] rule = "FFR309" path = "src/sqlbuild/compiler/planner/__init__.py" diff --git a/python-boundary.toml b/python-boundary.toml index 1a723fac71..99d850d278 100644 --- a/python-boundary.toml +++ b/python-boundary.toml @@ -2,7 +2,7 @@ # migration work that moves to Rust. `remaining_lines` must equal the current count of that work: # a PR that moves code lowers it, and adding movable Python fails the check. # Native failures must surface as errors; `handled_native_errors` lists the reviewed exceptions. -remaining_lines = 9602 +remaining_lines = 9601 [[stays]] path = "executor" @@ -434,16 +434,6 @@ path = "compiler/pipeline/_helpers/materializations.py" reason = "user-code" note = "Imports user materialization modules." -[[stays]] -path = "compiler/frontier/_helpers/code_identity.py" -reason = "user-code" -note = "Fingerprints the loaded Python modules behind cached compiler and rule results." - -[[stays]] -path = "compiler/frontier/main/compiled_code_identity.py" -reason = "user-code" -note = "Entry point for the loaded-code identity of cached compiler and rule results." - [[stays]] path = "compiler/macro_bridge/main" reason = "user-code" @@ -489,11 +479,6 @@ path = "compiler/auditing/main/get_audit_factory_definition.py" reason = "public-api" note = "Reads the declaration attached to a project audit factory." -[[stays]] -path = "compiler/auditing/main/_builtins.py" -reason = "public-api" -note = "Built-in audit factories written against the public audit factory API." - [[stays]] path = "compiler/sql_analysis/main/import_polyglot.py" reason = "public-api" @@ -509,11 +494,6 @@ path = "compiler/sql_analysis/_helpers/polyglot.py" reason = "public-api" note = "Wraps the Polyglot Python package with complexity guards." -[[stays]] -path = "compiler/profiling" -reason = "public-api" -note = "Compile timing collectors built on Python context managers and clocks." - [[stays]] path = "assets.py" reason = "public-api" diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index c41205c948..d118a64e3f 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -5172,8 +5172,7 @@ class SqlTestScanCache: SQL_TEST_SCAN_STORE_FILE_NAME: str def content_digest(parts: list[str]) -> str: ... -def installed_code_identity(facts: list[str], package_root: str | None) -> str: ... -def compiled_code_identity(installed: str, build_identity: str) -> str: ... +def compiled_code_identity() -> str: ... def fingerprint_project_files(root: str, excluded_files: list[str]) -> str: ... def digest_files(paths: list[str]) -> list[str | None]: ... diff --git a/src/sqlbuild/compiler/auditing/main/_builtins.py b/src/sqlbuild/compiler/auditing/main/_builtins.py deleted file mode 100644 index db5409dd37..0000000000 --- a/src/sqlbuild/compiler/auditing/main/_builtins.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Built-in generic audit definitions.""" - -from __future__ import annotations - -from pathlib import Path -from textwrap import dedent - -from sqlbuild.compiler.discovery.models import DiscoveredAuditBlock, DiscoveredAuditFile - - -def builtin_audit_definitions() -> tuple[DiscoveredAuditFile, ...]: - """Every built-in generic audit definition, which project definitions override by name.""" - - definitions: dict[str, tuple[DiscoveredAuditFile, DiscoveredAuditBlock]] = { - "accepted_values": _build_audit_definition( - name="accepted_values", - sql_body=dedent( - """\ - SELECT @column - FROM @relation - WHERE @column IS NOT NULL - AND @column NOT IN (@'values') - """ - ).strip(), - ), - "not_null": _build_audit_definition( - name="not_null", - sql_body=dedent( - """\ - SELECT @column - FROM @relation - WHERE @column IS NULL - """ - ).strip(), - ), - "relationships": _build_audit_definition( - name="relationships", - sql_body=dedent( - """\ - SELECT @column - FROM @relation - WHERE @column IS NOT NULL - AND @column NOT IN ( - SELECT @field - FROM @to - WHERE @field IS NOT NULL - ) - """ - ).strip(), - ), - "unique": _build_audit_definition( - name="unique", - sql_body=dedent( - """\ - SELECT @column, COUNT(*) AS duplicate_count - FROM @relation - WHERE @column IS NOT NULL - GROUP BY @column - HAVING COUNT(*) > 1 - """ - ).strip(), - ), - } - return tuple(audit_file for audit_file, _ in definitions.values()) - - -def _build_audit_definition( - *, name: str, sql_body: str -) -> tuple[DiscoveredAuditFile, DiscoveredAuditBlock]: - relative_path: Path = Path("audits/generic") / f"{name}.sql" - contents: str = f"AUDIT ();\n\n{sql_body}\n" - block: DiscoveredAuditBlock = DiscoveredAuditBlock( - audit_index=0, - header_values={}, - sql_body=sql_body, - ) - audit_file: DiscoveredAuditFile = DiscoveredAuditFile( - file_path=Path("") / relative_path, - relative_path=relative_path, - contents=contents, - blocks=(block,), - ) - return audit_file, block diff --git a/src/sqlbuild/compiler/frontier/_helpers/__init__.py b/src/sqlbuild/compiler/frontier/_helpers/__init__.py deleted file mode 100644 index e69de29bb2..0000000000 diff --git a/src/sqlbuild/compiler/frontier/_helpers/code_identity.py b/src/sqlbuild/compiler/frontier/_helpers/code_identity.py deleted file mode 100644 index 6886e97c61..0000000000 --- a/src/sqlbuild/compiler/frontier/_helpers/code_identity.py +++ /dev/null @@ -1,62 +0,0 @@ -"""Identity of the installed SQLBuild code that produces cached compiler and rule results.""" - -from __future__ import annotations - -import json -import platform -from functools import cache -from importlib.metadata import PackageNotFoundError, distribution -from pathlib import Path - -import sqlbuild._native as _native - -_PACKAGE_NAME: str = "sqlbuild" -_POLYGLOT_PACKAGE_NAME: str = "polyglot-sql-chio" -_YAML_PACKAGE_NAME: str = "PyYAML" -_DIRECT_URL_FILE: str = "direct_url.json" -_PACKAGE_ROOT_PARENT_DEPTH: int = 3 - - -@cache -def installed_code_identity() -> str: - """Return a digest that changes with every released or locally edited code version.""" - - facts: list[str] = [ - _package_version(_PACKAGE_NAME), - _package_version(_POLYGLOT_PACKAGE_NAME), - _package_version(_YAML_PACKAGE_NAME), - platform.python_implementation(), - ".".join(platform.python_version_tuple()[:2]), - ] - package_root: Path | None = ( - Path(__file__).resolve().parents[_PACKAGE_ROOT_PARENT_DEPTH] - if _is_editable_install() - else None - ) - return _native.installed_code_identity( - facts, None if package_root is None else str(package_root) - ) - - -def _package_version(package: str) -> str: - try: - return distribution(package).version - except PackageNotFoundError: - return "unknown" - - -def _is_editable_install() -> bool: - try: - direct_url: str | None = distribution(_PACKAGE_NAME).read_text(_DIRECT_URL_FILE) - except PackageNotFoundError: - return True - if direct_url is None: - return False - try: - payload: object = json.loads(direct_url) - except ValueError: - return True - if not isinstance(payload, dict): - return True - dir_info: object = payload.get("dir_info") - return isinstance(dir_info, dict) and dir_info.get("editable") is True diff --git a/src/sqlbuild/compiler/frontier/main/__init__.py b/src/sqlbuild/compiler/frontier/main/__init__.py deleted file mode 100644 index 40d7efab54..0000000000 --- a/src/sqlbuild/compiler/frontier/main/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Public compiler frontier entry points.""" diff --git a/src/sqlbuild/compiler/frontier/main/compiled_code_identity.py b/src/sqlbuild/compiler/frontier/main/compiled_code_identity.py deleted file mode 100644 index 4f4bcb6999..0000000000 --- a/src/sqlbuild/compiler/frontier/main/compiled_code_identity.py +++ /dev/null @@ -1,10 +0,0 @@ -"""Identity of the installed and native code behind cached compiler and rule results.""" - -import sqlbuild._native as _native -from sqlbuild.compiler.frontier._helpers.code_identity import installed_code_identity - - -def compiled_code_identity() -> str: - """Return the digest that changes with every released, locally edited, or rebuilt version.""" - - return _native.compiled_code_identity(installed_code_identity(), _native.BUILD_IDENTITY) diff --git a/src/sqlbuild/compiler/profiling/__init__.py b/src/sqlbuild/compiler/profiling/__init__.py deleted file mode 100644 index ce4ecbb6ec..0000000000 --- a/src/sqlbuild/compiler/profiling/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Compile profiling support.""" diff --git a/src/sqlbuild/compiler/profiling/_helpers/__init__.py b/src/sqlbuild/compiler/profiling/_helpers/__init__.py deleted file mode 100644 index 64fb62aa62..0000000000 --- a/src/sqlbuild/compiler/profiling/_helpers/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Native model analysis session helpers.""" diff --git a/src/sqlbuild/compiler/profiling/_helpers/timing.py b/src/sqlbuild/compiler/profiling/_helpers/timing.py deleted file mode 100644 index 42ccefd697..0000000000 --- a/src/sqlbuild/compiler/profiling/_helpers/timing.py +++ /dev/null @@ -1,23 +0,0 @@ -"""Compile phase timing by a chosen clock.""" - -from collections.abc import Callable, Iterator -from contextlib import contextmanager - -from sqlbuild.compiler.profiling.classes.context import CompileTimingContext -from sqlbuild.compiler.profiling.models import CompileTimingCollector -from sqlbuild.compiler.profiling.types import CompileTimingPhase - - -@contextmanager -def recorded_timing(*, phase: CompileTimingPhase, clock: Callable[[], int]) -> Iterator[None]: - """Record an operation by `clock` when a compile timing collector is active.""" - - collector: CompileTimingCollector | None = CompileTimingContext.active.get() - if collector is None: - yield - return - start_ns: int = clock() - try: - yield - finally: - collector.add(phase=phase, elapsed_ns=clock() - start_ns) diff --git a/src/sqlbuild/compiler/profiling/classes/__init__.py b/src/sqlbuild/compiler/profiling/classes/__init__.py deleted file mode 100644 index 666f0d4b99..0000000000 --- a/src/sqlbuild/compiler/profiling/classes/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Compile profiling runtime classes.""" diff --git a/src/sqlbuild/compiler/profiling/classes/context.py b/src/sqlbuild/compiler/profiling/classes/context.py deleted file mode 100644 index ab4ad2ffdd..0000000000 --- a/src/sqlbuild/compiler/profiling/classes/context.py +++ /dev/null @@ -1,14 +0,0 @@ -"""Active compile timing context.""" - -from contextvars import ContextVar -from typing import ClassVar - -from sqlbuild.compiler.profiling.models import CompileTimingCollector - - -class CompileTimingContext: - """Own the invocation-local active compile timing collector.""" - - active: ClassVar[ContextVar[CompileTimingCollector | None]] = ContextVar( - "active_compile_timings", default=None - ) diff --git a/src/sqlbuild/compiler/profiling/constants.py b/src/sqlbuild/compiler/profiling/constants.py deleted file mode 100644 index 7cf606a237..0000000000 --- a/src/sqlbuild/compiler/profiling/constants.py +++ /dev/null @@ -1,28 +0,0 @@ -"""Compile profiling constants.""" - -from sqlbuild.compiler.profiling.types import CompileMetric, CompileTimingPhase - -COMPILE_TIMING_PHASES: tuple[CompileTimingPhase, ...] = ( - "attachment_ms", - "model_analysis_ms", - "model_analysis_cpu_ms", - "analysis_preparation_ms", - "analysis_native_ms", - "analysis_projection_ms", - "binding_validation_ms", - "test_input_compile_ms", - "test_planning_ms", - "comparison_render_ms", - "cache_publication_ms", - "physical_write_ms", - "stale_traversal_ms", -) - -COMPILE_METRICS: tuple[CompileMetric, ...] = ( - "analysis_batch_cache_hits", - "analysis_entry_cache_hits", - "analysis_cache_misses", - "analysis_cache_bypasses", - "sql_test_scan_cache_hits", - "sql_test_scan_cache_misses", -) diff --git a/src/sqlbuild/compiler/profiling/main/__init__.py b/src/sqlbuild/compiler/profiling/main/__init__.py deleted file mode 100644 index 14f37a514d..0000000000 --- a/src/sqlbuild/compiler/profiling/main/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Public compile profiling entry points.""" diff --git a/src/sqlbuild/compiler/profiling/main/_collect.py b/src/sqlbuild/compiler/profiling/main/_collect.py deleted file mode 100644 index 50bb8290c6..0000000000 --- a/src/sqlbuild/compiler/profiling/main/_collect.py +++ /dev/null @@ -1,20 +0,0 @@ -"""Compile timing collection entry point.""" - -from collections.abc import Iterator -from contextlib import contextmanager -from contextvars import Token - -from sqlbuild.compiler.profiling.classes.context import CompileTimingContext -from sqlbuild.compiler.profiling.models import CompileTimingCollector - - -@contextmanager -def collect_compile_timings() -> Iterator[CompileTimingCollector]: - """Collect detailed timings for one compile command invocation.""" - - collector: CompileTimingCollector = CompileTimingCollector() - token: Token[CompileTimingCollector | None] = CompileTimingContext.active.set(collector) - try: - yield collector - finally: - CompileTimingContext.active.reset(token) diff --git a/src/sqlbuild/compiler/profiling/main/_metric.py b/src/sqlbuild/compiler/profiling/main/_metric.py deleted file mode 100644 index 17c7c6367a..0000000000 --- a/src/sqlbuild/compiler/profiling/main/_metric.py +++ /dev/null @@ -1,13 +0,0 @@ -"""Internal compile metric recording entry point.""" - -from sqlbuild.compiler.profiling.classes.context import CompileTimingContext -from sqlbuild.compiler.profiling.models import CompileTimingCollector -from sqlbuild.compiler.profiling.types import CompileMetric - - -def record_compile_metric(*, metric: CompileMetric, value: int) -> None: - """Record one non-negative metric when a compile collector is active.""" - - collector: CompileTimingCollector | None = CompileTimingContext.active.get() - if collector is not None: - collector.add_metric(metric=metric, value=max(0, value)) diff --git a/src/sqlbuild/compiler/profiling/main/_record.py b/src/sqlbuild/compiler/profiling/main/_record.py deleted file mode 100644 index 4053321d80..0000000000 --- a/src/sqlbuild/compiler/profiling/main/_record.py +++ /dev/null @@ -1,13 +0,0 @@ -"""Compile phase recording entry point.""" - -import time -from contextlib import AbstractContextManager - -from sqlbuild.compiler.profiling._helpers.timing import recorded_timing -from sqlbuild.compiler.profiling.types import CompileTimingPhase - - -def record_compile_timing(phase: CompileTimingPhase) -> AbstractContextManager[None]: - """Record an operation when a compile timing collector is active.""" - - return recorded_timing(phase=phase, clock=time.perf_counter_ns) diff --git a/src/sqlbuild/compiler/profiling/main/_record_cpu.py b/src/sqlbuild/compiler/profiling/main/_record_cpu.py deleted file mode 100644 index a0a15e0ef7..0000000000 --- a/src/sqlbuild/compiler/profiling/main/_record_cpu.py +++ /dev/null @@ -1,13 +0,0 @@ -"""Compile phase CPU recording entry point.""" - -import time -from contextlib import AbstractContextManager - -from sqlbuild.compiler.profiling._helpers.timing import recorded_timing -from sqlbuild.compiler.profiling.types import CompileTimingPhase - - -def record_compile_cpu_timing(phase: CompileTimingPhase) -> AbstractContextManager[None]: - """Record the process CPU time an operation spans, including its worker threads.""" - - return recorded_timing(phase=phase, clock=time.process_time_ns) diff --git a/src/sqlbuild/compiler/profiling/models.py b/src/sqlbuild/compiler/profiling/models.py deleted file mode 100644 index fd6ad7e715..0000000000 --- a/src/sqlbuild/compiler/profiling/models.py +++ /dev/null @@ -1,31 +0,0 @@ -"""Compile profiling models.""" - -from dataclasses import dataclass, field - -from sqlbuild.compiler.profiling.constants import COMPILE_METRICS, COMPILE_TIMING_PHASES -from sqlbuild.compiler.profiling.types import CompileMetric, CompileTimingPhase - - -@dataclass(frozen=True) -class CompileTimingCollector: - """Accumulate elapsed nanoseconds for nested compile operations.""" - - elapsed_ns: dict[CompileTimingPhase, int] = field(default_factory=dict) - metrics: dict[CompileMetric, int] = field(default_factory=dict) - - def add(self, *, phase: CompileTimingPhase, elapsed_ns: int) -> None: - """Add one measured operation to a compile phase.""" - - self.elapsed_ns[phase] = self.elapsed_ns.get(phase, 0) + elapsed_ns - - def add_metric(self, *, metric: CompileMetric, value: int) -> None: - """Add a non-negative compile metric to the current invocation.""" - - self.metrics[metric] = self.metrics.get(metric, 0) + value - - def as_milliseconds(self) -> dict[str, int]: - """Return every supported phase in stable presentation order.""" - - return { - phase: self.elapsed_ns.get(phase, 0) // 1_000_000 for phase in COMPILE_TIMING_PHASES - } | {metric: self.metrics.get(metric, 0) for metric in COMPILE_METRICS} diff --git a/src/sqlbuild/rule_engine/_helpers/engine/custom_dispatch.py b/src/sqlbuild/rule_engine/_helpers/engine/custom_dispatch.py index ae211e78c6..1f5b2a093e 100644 --- a/src/sqlbuild/rule_engine/_helpers/engine/custom_dispatch.py +++ b/src/sqlbuild/rule_engine/_helpers/engine/custom_dispatch.py @@ -16,7 +16,6 @@ selected_rule_codes, ) from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.compiler.frontier.main.compiled_code_identity import compiled_code_identity from sqlbuild.rule_engine._helpers.engine.custom_payloads import custom_rule_payloads from sqlbuild.rule_engine._helpers.engine.custom_rules import evaluate_custom_rules_cached from sqlbuild.rule_engine.constants import TYPE_PROOF_RULE_CODES @@ -148,7 +147,7 @@ def start_custom_rules( if not custom_rules: return None if config.cache.enabled: - _ = compiled_code_identity() + _ = _native.compiled_code_identity() return executor.submit( evaluate_custom_rules_cached, project=project, diff --git a/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py b/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py index 5acdaa97b9..4f6af8389d 100644 --- a/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py +++ b/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py @@ -9,6 +9,7 @@ from pathlib import Path from typing import Any +import sqlbuild._native as _native from sqlbuild._native import ( CustomRuleCache, compiled_project_record, @@ -19,7 +20,6 @@ ) from sqlbuild.compiler.compile.models import CompiledProject from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.compiler.frontier.main.compiled_code_identity import compiled_code_identity from sqlbuild.compiler.scopes.models import ScopeIndex from sqlbuild.rule_engine._helpers.engine.fact_replay import full_fact_keys, is_tree_fact from sqlbuild.rule_engine._helpers.host.custom_host_pool import run_custom_hosts @@ -130,7 +130,7 @@ def _rule_identity( *, rule: Rule, config: RulesConfig, project_dir: Path, dialect: str, fingerprint: str | None ) -> str: return custom_rule_identity( - code_identity=compiled_code_identity(), + code_identity=_native.compiled_code_identity(), runtime_version=CUSTOM_HOST_RUNTIME_VERSION, code=rule.code, family=rule.family, diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/compile/helpers.py b/tests/e2e/src/sqlbuild/cli/commands/main/compile/helpers.py index 0b756cad53..f27228d539 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/compile/helpers.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/compile/helpers.py @@ -34,7 +34,6 @@ import sqlbuild.adapter.type_system._helpers.type_normalization as type_normalization import sqlbuild.cli.compile_reuse._helpers.provider_settings as provider_settings import sqlbuild.compiler.compile._helpers.macro_bridge.call_store as call_store_module -import sqlbuild.compiler.frontier.main.compiled_code_identity as compiled_code_identity_module import sqlbuild.compiler.macro_bridge.classes.macro_bridge as macro_bridge_class from scripts.cold_compile_performance.main.read_compile_measurement import read_compile_measurement from scripts.cold_compile_performance.main.semantic_compile_fingerprint import ( @@ -4120,9 +4119,7 @@ def corrupt_sql_test_scan_store(root: Path, _monkeypatch: pytest.MonkeyPatch) -> def edit_installed_code(_root: Path, monkeypatch: pytest.MonkeyPatch) -> None: """Pretend the installed Python code changed, as a local edit of an editable install does.""" - monkeypatch.setattr( - compiled_code_identity_module, "installed_code_identity", lambda: "edited-python-code" - ) + monkeypatch.setattr(native_module, "compiled_code_identity", lambda: "edited-python-code") RETIRED_RENDER_FILES: tuple[str, ...] = ( diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/run/helpers.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/run/helpers.py index 4a74d73192..20cdb0c034 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/run/helpers.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/run/helpers.py @@ -4,7 +4,6 @@ from sqlbuild import _native from sqlbuild._native import NativeRulesRequest -from sqlbuild.compiler.frontier.main import compiled_code_identity as code_identity from sqlbuild.rule_engine._helpers.engine import custom_dispatch as native from sqlbuild.rule_engine.models import NativeRulesEvaluation @@ -39,5 +38,5 @@ def recording(request: NativeRulesRequest) -> NativeRulesEvaluation: def rebuild_identity(*, monkeypatch: pytest.MonkeyPatch) -> None: """Pretend the compiled code changed, as after installing another build.""" - identity: str = code_identity.compiled_code_identity() - monkeypatch.setattr(code_identity, "compiled_code_identity", lambda: f"{identity}-rebuilt") + identity: str = _native.compiled_code_identity() + monkeypatch.setattr(_native, "compiled_code_identity", lambda: f"{identity}-rebuilt") From 23b1114f0806d18d153169b05c88c117bb8d58d1 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 15 Oct 2026 09:43:40 +0100 Subject: [PATCH 02/12] refactor: run dbt subprocesses, event streams and project init natively --- .../src/dbt_process/_helpers/exit_codes.rs | 22 + .../src/dbt_process/_helpers/mod.rs | 2 + .../src/dbt_process/_helpers/text.rs | 8 + .../src/dbt_process/constants.rs | 13 + .../main/dbt_process_returncode.rs | 11 + .../sqlbuild-core/src/dbt_process/main/mod.rs | 5 + .../dbt_process/main/read_dbt_output_line.rs | 14 + .../src/dbt_process/main/run_dbt_process.rs | 30 ++ .../main/spawn_dbt_json_process.rs | 42 ++ .../src/dbt_process/main/stop_dbt_process.rs | 33 ++ crates/sqlbuild-core/src/dbt_process/mod.rs | 9 + .../sqlbuild-core/src/dbt_process/models.rs | 19 + crates/sqlbuild-core/src/dbt_process/tests.rs | 6 + .../src/dbt_process/tests/helpers.rs | 12 + .../dbt_process/tests/test_dbt_processes.rs | 151 ++++++ .../src/dbt_process/tests/test_types.rs | 28 ++ crates/sqlbuild-core/src/lib.rs | 1 + .../_helpers/dbt_command/auto_init.rs | 4 +- .../src/bindings/_helpers/dbt_command/init.rs | 6 +- .../_helpers/dbt_command/init_project.rs | 369 ++++++++++++++ .../src/bindings/_helpers/dbt_command/mod.rs | 1 + .../_helpers/dbt_command/registration.rs | 5 +- .../src/bindings/_helpers/dbt_interop/cli.rs | 11 +- .../_helpers/dbt_interop/event_stream.rs | 437 +++++++++++++++++ .../src/bindings/_helpers/dbt_interop/mod.rs | 2 + .../bindings/_helpers/dbt_interop/pipeline.rs | 6 +- .../_helpers/dbt_interop/pipeline_prologue.rs | 6 +- .../_helpers/dbt_interop/registration.rs | 2 + .../bindings/_helpers/dbt_interop/runner.rs | 296 +++++++++++ .../_helpers/dbt_interop/runtime_execution.rs | 6 +- fensu.toml | 5 + python-boundary.toml | 22 +- src/sqlbuild/_native.pyi | 53 +- .../integrations/dbt/_helpers/profile/init.py | 207 -------- .../dbt/_helpers/runtime/event_stream.py | 405 --------------- .../integrations/dbt/classes/__init__.py | 1 - .../integrations/dbt/classes/dbt_runner.py | 125 ----- .../integrations/dbt/main/runtime/__init__.py | 1 - .../runtime/_execute_json_event_stream.py | 37 -- .../sqlbuild/integrations/dbt/test_graph.py | 2 +- .../integrations/dbt/test_manifest.py | 3 +- .../sqlbuild/integrations/dbt/test_runner.py | 2 +- .../cli/commands/main/dbt/test_dbt.py | 4 +- .../cli/commands/main/dbt/test_init.py | 2 +- .../sqlbuild/integrations/dbt/_test_types.py | 16 + .../src/sqlbuild/integrations/dbt/helpers.py | 33 +- .../integrations/dbt/profile/test_init.py | 2 +- .../sqlbuild/integrations/dbt/test_debug.py | 3 +- .../integrations/dbt/test_event_stream.py | 463 ++---------------- .../sqlbuild/integrations/dbt/test_execute.py | 4 +- .../dbt/test_plan_orchestration.py | 2 +- .../sqlbuild/integrations/dbt/test_runner.py | 65 ++- 52 files changed, 1756 insertions(+), 1258 deletions(-) create mode 100644 crates/sqlbuild-core/src/dbt_process/_helpers/exit_codes.rs create mode 100644 crates/sqlbuild-core/src/dbt_process/_helpers/mod.rs create mode 100644 crates/sqlbuild-core/src/dbt_process/_helpers/text.rs create mode 100644 crates/sqlbuild-core/src/dbt_process/constants.rs create mode 100644 crates/sqlbuild-core/src/dbt_process/main/dbt_process_returncode.rs create mode 100644 crates/sqlbuild-core/src/dbt_process/main/mod.rs create mode 100644 crates/sqlbuild-core/src/dbt_process/main/read_dbt_output_line.rs create mode 100644 crates/sqlbuild-core/src/dbt_process/main/run_dbt_process.rs create mode 100644 crates/sqlbuild-core/src/dbt_process/main/spawn_dbt_json_process.rs create mode 100644 crates/sqlbuild-core/src/dbt_process/main/stop_dbt_process.rs create mode 100644 crates/sqlbuild-core/src/dbt_process/mod.rs create mode 100644 crates/sqlbuild-core/src/dbt_process/models.rs create mode 100644 crates/sqlbuild-core/src/dbt_process/tests.rs create mode 100644 crates/sqlbuild-core/src/dbt_process/tests/helpers.rs create mode 100644 crates/sqlbuild-core/src/dbt_process/tests/test_dbt_processes.rs create mode 100644 crates/sqlbuild-core/src/dbt_process/tests/test_types.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/dbt_command/init_project.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/event_stream.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/runner.rs delete mode 100644 src/sqlbuild/integrations/dbt/_helpers/profile/init.py delete mode 100644 src/sqlbuild/integrations/dbt/_helpers/runtime/event_stream.py delete mode 100644 src/sqlbuild/integrations/dbt/classes/__init__.py delete mode 100644 src/sqlbuild/integrations/dbt/classes/dbt_runner.py delete mode 100644 src/sqlbuild/integrations/dbt/main/runtime/__init__.py delete mode 100644 src/sqlbuild/integrations/dbt/main/runtime/_execute_json_event_stream.py diff --git a/crates/sqlbuild-core/src/dbt_process/_helpers/exit_codes.rs b/crates/sqlbuild-core/src/dbt_process/_helpers/exit_codes.rs new file mode 100644 index 0000000000..ed4d0d1e8e --- /dev/null +++ b/crates/sqlbuild-core/src/dbt_process/_helpers/exit_codes.rs @@ -0,0 +1,22 @@ +//! `subprocess` return codes from exit statuses. + +use std::process::ExitStatus; + +/// The exit code, or the negated signal number of a signalled process. +pub(crate) fn returncode(status: ExitStatus) -> i64 { + if let Some(code) = status.code() { + return i64::from(code); + } + signal_returncode(status) +} + +#[cfg(unix)] +fn signal_returncode(status: ExitStatus) -> i64 { + use std::os::unix::process::ExitStatusExt; + status.signal().map_or(-1, |signal| -i64::from(signal)) +} + +#[cfg(not(unix))] +fn signal_returncode(_status: ExitStatus) -> i64 { + -1 +} diff --git a/crates/sqlbuild-core/src/dbt_process/_helpers/mod.rs b/crates/sqlbuild-core/src/dbt_process/_helpers/mod.rs new file mode 100644 index 0000000000..b2df644c67 --- /dev/null +++ b/crates/sqlbuild-core/src/dbt_process/_helpers/mod.rs @@ -0,0 +1,2 @@ +pub(crate) mod exit_codes; +pub(crate) mod text; diff --git a/crates/sqlbuild-core/src/dbt_process/_helpers/text.rs b/crates/sqlbuild-core/src/dbt_process/_helpers/text.rs new file mode 100644 index 0000000000..1c84a8701f --- /dev/null +++ b/crates/sqlbuild-core/src/dbt_process/_helpers/text.rs @@ -0,0 +1,8 @@ +//! Python text-mode decoding of subprocess output. + +/// Lossy UTF-8 with universal newlines: `\r\n` and lone `\r` become `\n`. +pub(crate) fn universal_text(bytes: &[u8]) -> String { + String::from_utf8_lossy(bytes) + .replace("\r\n", "\n") + .replace('\r', "\n") +} diff --git a/crates/sqlbuild-core/src/dbt_process/constants.rs b/crates/sqlbuild-core/src/dbt_process/constants.rs new file mode 100644 index 0000000000..c8e13d015e --- /dev/null +++ b/crates/sqlbuild-core/src/dbt_process/constants.rs @@ -0,0 +1,13 @@ +//! Environment and timing of dbt subprocesses. + +use std::time::Duration; + +/// Environment a JSON log stream adds to the inherited one. +pub const DBT_JSON_ENVIRONMENT: [(&str, &str); 2] = + [("DBT_LOG_FORMAT", "json"), ("PYTHONUNBUFFERED", "1")]; +/// Variables pointing dbt's artifacts and logs at the target folder. +pub const DBT_TARGET_PATH_VARIABLES: [&str; 2] = ["DBT_TARGET_PATH", "DBT_LOG_PATH"]; +/// How long a stopping dbt process may take after each terminate or kill. +pub const DBT_STOP_WAIT: Duration = Duration::from_secs(1); +/// How often a stopping dbt process is polled. +pub(crate) const DBT_STOP_POLL: Duration = Duration::from_millis(10); diff --git a/crates/sqlbuild-core/src/dbt_process/main/dbt_process_returncode.rs b/crates/sqlbuild-core/src/dbt_process/main/dbt_process_returncode.rs new file mode 100644 index 0000000000..7049b4d85b --- /dev/null +++ b/crates/sqlbuild-core/src/dbt_process/main/dbt_process_returncode.rs @@ -0,0 +1,11 @@ +//! Wait for a dbt process and report its `subprocess` return code. + +use std::io; +use std::process::Child; + +use crate::dbt_process::_helpers::exit_codes::returncode; + +/// The exit code, or the negated signal number, once the process has ended. +pub fn dbt_process_returncode(child: &mut Child) -> io::Result { + Ok(returncode(child.wait()?)) +} diff --git a/crates/sqlbuild-core/src/dbt_process/main/mod.rs b/crates/sqlbuild-core/src/dbt_process/main/mod.rs new file mode 100644 index 0000000000..f688d594dd --- /dev/null +++ b/crates/sqlbuild-core/src/dbt_process/main/mod.rs @@ -0,0 +1,5 @@ +pub mod dbt_process_returncode; +pub mod read_dbt_output_line; +pub mod run_dbt_process; +pub mod spawn_dbt_json_process; +pub mod stop_dbt_process; diff --git a/crates/sqlbuild-core/src/dbt_process/main/read_dbt_output_line.rs b/crates/sqlbuild-core/src/dbt_process/main/read_dbt_output_line.rs new file mode 100644 index 0000000000..1b97f4daf6 --- /dev/null +++ b/crates/sqlbuild-core/src/dbt_process/main/read_dbt_output_line.rs @@ -0,0 +1,14 @@ +//! The next text line of a dbt JSON log stream. + +use std::io::{self, BufRead, BufReader, PipeReader}; + +use crate::dbt_process::_helpers::text::universal_text; + +/// The next line with its newline, decoded as Python text mode would; `None` at the end. +pub fn read_dbt_output_line(output: &mut BufReader) -> io::Result> { + let mut line: Vec = Vec::new(); + if output.read_until(b'\n', &mut line)? == 0 { + return Ok(None); + } + Ok(Some(universal_text(&line))) +} diff --git a/crates/sqlbuild-core/src/dbt_process/main/run_dbt_process.rs b/crates/sqlbuild-core/src/dbt_process/main/run_dbt_process.rs new file mode 100644 index 0000000000..ee0ae6824d --- /dev/null +++ b/crates/sqlbuild-core/src/dbt_process/main/run_dbt_process.rs @@ -0,0 +1,30 @@ +//! Run one dbt command to completion with its output captured as text. + +use std::io; +use std::path::Path; +use std::process::{Command, Output}; + +use crate::dbt_process::_helpers::exit_codes::returncode; +use crate::dbt_process::_helpers::text::universal_text; +use crate::dbt_process::models::DbtProcessOutput; + +/// `subprocess.run(argv, cwd=cwd, capture_output=True, text=True)`. +pub fn run_dbt_process(argv: &[String], cwd: Option<&Path>) -> io::Result { + let Some((program, arguments)) = argv.split_first() else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "empty dbt argv", + )); + }; + let mut command: Command = Command::new(program); + command.args(arguments); + if let Some(directory) = cwd { + command.current_dir(directory); + } + let output: Output = command.output()?; + Ok(DbtProcessOutput { + returncode: returncode(output.status), + stdout: universal_text(&output.stdout), + stderr: universal_text(&output.stderr), + }) +} diff --git a/crates/sqlbuild-core/src/dbt_process/main/spawn_dbt_json_process.rs b/crates/sqlbuild-core/src/dbt_process/main/spawn_dbt_json_process.rs new file mode 100644 index 0000000000..3f97ec4bf7 --- /dev/null +++ b/crates/sqlbuild-core/src/dbt_process/main/spawn_dbt_json_process.rs @@ -0,0 +1,42 @@ +//! Start dbt with JSON logs, its stdout and stderr merged into one line stream. + +use std::io::{self, BufReader, PipeReader, PipeWriter}; +use std::path::Path; +use std::process::{Command, Stdio}; + +use crate::dbt_process::constants::{DBT_JSON_ENVIRONMENT, DBT_TARGET_PATH_VARIABLES}; +use crate::dbt_process::models::DbtJsonProcess; + +/// Spawn `argv` in `cwd` with JSON logging and, when given, artifacts and logs below `target_path`. +pub fn spawn_dbt_json_process( + argv: &[String], + (cwd, target_path): (Option<&Path>, Option<&Path>), +) -> io::Result { + let Some((program, arguments)) = argv.split_first() else { + return Err(io::Error::new( + io::ErrorKind::InvalidInput, + "empty dbt argv", + )); + }; + let (reader, writer): (PipeReader, PipeWriter) = io::pipe()?; + let mut command: Command = Command::new(program); + command + .args(arguments) + .envs(DBT_JSON_ENVIRONMENT) + .stdout(Stdio::from(writer.try_clone()?)) + .stderr(Stdio::from(writer)); + if let Some(target) = target_path { + for variable in DBT_TARGET_PATH_VARIABLES { + command.env(variable, target); + } + } + if let Some(directory) = cwd { + command.current_dir(directory); + } + let child = command.spawn()?; + drop(command); + Ok(DbtJsonProcess { + child, + output: BufReader::new(reader), + }) +} diff --git a/crates/sqlbuild-core/src/dbt_process/main/stop_dbt_process.rs b/crates/sqlbuild-core/src/dbt_process/main/stop_dbt_process.rs new file mode 100644 index 0000000000..8072e6ddff --- /dev/null +++ b/crates/sqlbuild-core/src/dbt_process/main/stop_dbt_process.rs @@ -0,0 +1,33 @@ +//! Stop an unfinished dbt process: terminate, wait, then kill and wait again. + +use std::io; +use std::process::{Child, ExitStatus}; +use std::time::Instant; + +use crate::dbt_process::constants::{DBT_STOP_POLL, DBT_STOP_WAIT}; + +/// Terminate a running process with `terminate`, kill it after a second, and reap it. +pub fn stop_dbt_process(mut child: Child, terminate: impl FnOnce(u32)) -> Child { + if matches!(child.try_wait(), Ok(None)) { + terminate(child.id()); + } + if matches!(wait_briefly(&mut child), Ok(Some(_))) { + return child; + } + let _killed: io::Result<()> = child.kill(); + let _reaped: io::Result> = wait_briefly(&mut child); + child +} + +fn wait_briefly(child: &mut Child) -> io::Result> { + let deadline: Instant = Instant::now() + DBT_STOP_WAIT; + loop { + if let Some(status) = child.try_wait()? { + return Ok(Some(status)); + } + if Instant::now() >= deadline { + return Ok(None); + } + std::thread::sleep(DBT_STOP_POLL); + } +} diff --git a/crates/sqlbuild-core/src/dbt_process/mod.rs b/crates/sqlbuild-core/src/dbt_process/mod.rs new file mode 100644 index 0000000000..90a9dc9e65 --- /dev/null +++ b/crates/sqlbuild-core/src/dbt_process/mod.rs @@ -0,0 +1,9 @@ +//! dbt CLI subprocesses: captured commands and the JSON log stream with its shutdown. + +mod _helpers; +pub mod constants; +pub mod main; +pub mod models; + +#[cfg(test)] +mod tests; diff --git a/crates/sqlbuild-core/src/dbt_process/models.rs b/crates/sqlbuild-core/src/dbt_process/models.rs new file mode 100644 index 0000000000..432d321f64 --- /dev/null +++ b/crates/sqlbuild-core/src/dbt_process/models.rs @@ -0,0 +1,19 @@ +//! A finished captured dbt command and a running dbt JSON log stream. + +use std::io::{BufReader, PipeReader}; +use std::process::Child; + +/// Exit code (negative signal numbers for signalled processes) and decoded output. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct DbtProcessOutput { + pub returncode: i64, + pub stdout: String, + pub stderr: String, +} + +/// A running dbt process whose stdout and stderr share one line reader. +#[derive(Debug)] +pub struct DbtJsonProcess { + pub child: Child, + pub output: BufReader, +} diff --git a/crates/sqlbuild-core/src/dbt_process/tests.rs b/crates/sqlbuild-core/src/dbt_process/tests.rs new file mode 100644 index 0000000000..42f6cd1508 --- /dev/null +++ b/crates/sqlbuild-core/src/dbt_process/tests.rs @@ -0,0 +1,6 @@ +#[path = "tests/test_dbt_processes.rs"] +mod dbt_processes; +#[path = "tests/helpers.rs"] +mod helpers; +#[path = "tests/test_types.rs"] +mod test_types; diff --git a/crates/sqlbuild-core/src/dbt_process/tests/helpers.rs b/crates/sqlbuild-core/src/dbt_process/tests/helpers.rs new file mode 100644 index 0000000000..9d0cd66c7f --- /dev/null +++ b/crates/sqlbuild-core/src/dbt_process/tests/helpers.rs @@ -0,0 +1,12 @@ +use crate::dbt_process::main::read_dbt_output_line::read_dbt_output_line; +use crate::dbt_process::models::DbtJsonProcess; + +/// argv running `script` with `sh -c`. +pub(super) fn shell(script: &str) -> Vec { + vec!["sh".to_owned(), "-c".to_owned(), script.to_owned()] +} + +/// Every line a streamed process writes, until it closes its output. +pub(super) fn collected_lines(process: &mut DbtJsonProcess) -> Vec { + std::iter::from_fn(|| read_dbt_output_line(&mut process.output).expect("readable")).collect() +} diff --git a/crates/sqlbuild-core/src/dbt_process/tests/test_dbt_processes.rs b/crates/sqlbuild-core/src/dbt_process/tests/test_dbt_processes.rs new file mode 100644 index 0000000000..f085bb94de --- /dev/null +++ b/crates/sqlbuild-core/src/dbt_process/tests/test_dbt_processes.rs @@ -0,0 +1,151 @@ +#![cfg(unix)] + +use std::path::Path; + +use crate::dbt_process::main::dbt_process_returncode::dbt_process_returncode; +use crate::dbt_process::main::run_dbt_process::run_dbt_process; +use crate::dbt_process::main::spawn_dbt_json_process::spawn_dbt_json_process; +use crate::dbt_process::main::stop_dbt_process::stop_dbt_process; +use crate::dbt_process::tests::helpers::{collected_lines, shell}; +use crate::dbt_process::tests::test_types::{ + CapturedProcessTestCase, LaunchFailureTestCase, StopProcessTestCase, StreamedProcessTestCase, +}; + +#[test] +fn given_dbt_command_when_captured_then_output_and_exit_code_match_python_text_mode() { + let test_cases = [ + CapturedProcessTestCase { + description: "a clean exit keeps both streams with universal newlines", + script: "printf 'one\\r\\ntwo\\r'; printf 'warn\\n' >&2", + expected_returncode: 0, + expected_stdout: "one\ntwo\n", + expected_stderr: "warn\n", + }, + CapturedProcessTestCase { + description: "a failing exit reports its code", + script: "exit 3", + expected_returncode: 3, + expected_stdout: "", + expected_stderr: "", + }, + CapturedProcessTestCase { + description: "a signalled process reports the negated signal", + script: "kill -TERM $$", + expected_returncode: -15, + expected_stdout: "", + expected_stderr: "", + }, + ]; + for test_case in test_cases { + let output = run_dbt_process(&shell(test_case.script), None).expect("the shell runs"); + + assert_eq!( + output.returncode, test_case.expected_returncode, + "{}", + test_case.description + ); + assert_eq!( + output.stdout, test_case.expected_stdout, + "{}", + test_case.description + ); + assert_eq!( + output.stderr, test_case.expected_stderr, + "{}", + test_case.description + ); + } +} + +#[test] +fn given_missing_dbt_executable_when_run_then_launch_error_is_not_found() { + let test_cases = [LaunchFailureTestCase { + description: "a missing program cannot start", + program: "/nonexistent/dbt", + expected_kind: std::io::ErrorKind::NotFound, + }]; + for test_case in test_cases { + let error = run_dbt_process(&[test_case.program.to_owned()], None).expect_err("no binary"); + + assert_eq!( + error.kind(), + test_case.expected_kind, + "{}", + test_case.description + ); + } +} + +#[test] +fn given_dbt_json_process_when_streamed_then_lines_merge_both_streams_with_json_env() { + let test_cases = [ + StreamedProcessTestCase { + description: "stdout and stderr share one stream with the JSON environment", + script: "echo \"$DBT_LOG_FORMAT $PYTHONUNBUFFERED ${DBT_TARGET_PATH:-none}\"; echo err >&2; printf tail", + target_path: None, + expected_lines: vec!["json 1 none\n", "err\n", "tail"], + expected_returncode: 0, + }, + StreamedProcessTestCase { + description: "a target path sets the target and log folders", + script: "echo \"$DBT_TARGET_PATH $DBT_LOG_PATH\"; exit 4", + target_path: Some("/tmp/target"), + expected_lines: vec!["/tmp/target /tmp/target\n"], + expected_returncode: 4, + }, + ]; + for test_case in test_cases { + let mut process = spawn_dbt_json_process( + &shell(test_case.script), + (None, test_case.target_path.map(Path::new)), + ) + .expect("the shell starts"); + let lines: Vec = collected_lines(&mut process); + + assert_eq!(lines, test_case.expected_lines, "{}", test_case.description); + assert_eq!( + dbt_process_returncode(&mut process.child).expect("waitable"), + test_case.expected_returncode, + "{}", + test_case.description + ); + } +} + +#[test] +fn given_running_dbt_process_when_stopped_then_terminate_is_requested_and_child_is_reaped() { + let test_cases = [ + StopProcessTestCase { + description: "a running process is asked to terminate, then killed and reaped", + script: "exec sleep 30", + settle: std::time::Duration::ZERO, + expected_terminate_requests: 1, + }, + StopProcessTestCase { + description: "an exited process is only reaped", + script: "exit 0", + settle: std::time::Duration::from_millis(200), + expected_terminate_requests: 0, + }, + ]; + for test_case in test_cases { + let process = + spawn_dbt_json_process(&shell(test_case.script), (None, None)).expect("starts"); + std::thread::sleep(test_case.settle); + let mut requests: Vec = Vec::new(); + + let mut child = stop_dbt_process(process.child, |pid| requests.push(pid)); + + assert_eq!( + requests.len(), + test_case.expected_terminate_requests, + "{}", + test_case.description + ); + assert!( + child.try_wait().expect("waitable").is_some(), + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-core/src/dbt_process/tests/test_types.rs b/crates/sqlbuild-core/src/dbt_process/tests/test_types.rs new file mode 100644 index 0000000000..7628b66a42 --- /dev/null +++ b/crates/sqlbuild-core/src/dbt_process/tests/test_types.rs @@ -0,0 +1,28 @@ +pub(super) struct CapturedProcessTestCase { + pub(super) description: &'static str, + pub(super) script: &'static str, + pub(super) expected_returncode: i64, + pub(super) expected_stdout: &'static str, + pub(super) expected_stderr: &'static str, +} + +pub(super) struct StreamedProcessTestCase { + pub(super) description: &'static str, + pub(super) script: &'static str, + pub(super) target_path: Option<&'static str>, + pub(super) expected_lines: Vec<&'static str>, + pub(super) expected_returncode: i64, +} + +pub(super) struct LaunchFailureTestCase { + pub(super) description: &'static str, + pub(super) program: &'static str, + pub(super) expected_kind: std::io::ErrorKind, +} + +pub(super) struct StopProcessTestCase { + pub(super) description: &'static str, + pub(super) script: &'static str, + pub(super) settle: std::time::Duration, + pub(super) expected_terminate_requests: usize, +} diff --git a/crates/sqlbuild-core/src/lib.rs b/crates/sqlbuild-core/src/lib.rs index bbb65f2801..fce25f3c1c 100644 --- a/crates/sqlbuild-core/src/lib.rs +++ b/crates/sqlbuild-core/src/lib.rs @@ -20,6 +20,7 @@ pub mod dbt_command; pub mod dbt_events; pub mod dbt_graph; pub mod dbt_plan; +pub mod dbt_process; pub mod debug_report; pub mod format_command; pub mod freshness_text; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/auto_init.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/auto_init.rs index 7d505418ba..1631392a1a 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/auto_init.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/auto_init.rs @@ -14,9 +14,9 @@ use sqlbuild_core::dbt_command::models::ForwardedArg; use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported}; use crate::bindings::_helpers::dbt_command::support::{color_enabled, flag, path, stream}; +const NATIVE_MODULE: &str = "sqlbuild._native"; const DISCOVERY_CONSTANTS: &str = "sqlbuild.compiler.discovery.constants"; const PROGRESS_REPORTER_MODULE: &str = "sqlbuild._native"; -const PROFILE_INIT_MODULE: &str = "sqlbuild.integrations.dbt._helpers.profile.init"; const DBT_MODELS: &str = "sqlbuild.integrations.dbt.models"; fn project_config_found(directory: &Bound<'_, PyAny>) -> PyResult { @@ -122,7 +122,7 @@ pub(crate) fn ensure_twin_project<'py>( ], )?; let result: Bound<'py, PyAny> = call_with( - &imported(py, PROFILE_INIT_MODULE, "build_dbt_init_project")?, + &imported(py, NATIVE_MODULE, "build_dbt_init_project")?, &[("request", request)], )?; let project_file: String = result.getattr("project_file")?.str()?.to_string(); diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/init.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/init.rs index 2c21464486..8533136121 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/init.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/init.rs @@ -14,9 +14,9 @@ use sqlbuild_core::dbt_command::models::DbtInitSummary; use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported}; use crate::bindings::_helpers::dbt_command::support::{color_enabled, flag, path, stream}; +const NATIVE_MODULE: &str = "sqlbuild._native"; const ERRORS_MODULE: &str = "sqlbuild._native"; const PROGRESS_REPORTER_MODULE: &str = "sqlbuild._native"; -const PROFILE_INIT_MODULE: &str = "sqlbuild.integrations.dbt._helpers.profile.init"; const DBT_MODELS: &str = "sqlbuild.integrations.dbt.models"; fn missing_project(py: Python<'_>) -> PyErr { @@ -118,11 +118,11 @@ fn run_dbt_init_command(py: Python<'_>, request: &Bound<'_, PyAny>) -> PyResult< ], )?; call_with( - &imported(py, PROFILE_INIT_MODULE, "validate_dbt_init_project")?, + &imported(py, NATIVE_MODULE, "validate_dbt_init_project")?, &[("request", init_request.clone())], )?; let result: Bound<'_, PyAny> = call_with( - &imported(py, PROFILE_INIT_MODULE, "build_dbt_init_project")?, + &imported(py, NATIVE_MODULE, "build_dbt_init_project")?, &[("request", init_request)], )?; let style: CliStyle = CliStyle { use_color }; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/init_project.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/init_project.rs new file mode 100644 index 0000000000..407c1345a5 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/init_project.rs @@ -0,0 +1,369 @@ +//! Build or validate a SQLBuild project from a dbt project, its profile and `dbt debug`. + +use std::io; + +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::{ + Bound, PyAny, PyAnyMethods, PyErr, PyModule, PyModuleMethods, PyResult, Python, +}; +use pyo3::types::{PyDict, PyDictMethods, PyString, PyTuple}; +use pyo3::{IntoPyObjectExt, pyfunction, wrap_pyfunction}; +use sqlbuild_core::dbt_process::main::run_dbt_process::run_dbt_process; +use sqlbuild_core::dbt_process::models::DbtProcessOutput; +use sqlbuild_core::text::main::python_strip::python_strip; + +use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported}; +use crate::bindings::_helpers::dbt_interop::runner::launch_os_error; +use crate::bindings::types::CompilerDetach; + +const NATIVE_MODULE: &str = "sqlbuild._native"; +const DBT_MODELS: &str = "sqlbuild.integrations.dbt.models"; +const DBT_EXCEPTIONS: &str = "sqlbuild.integrations.dbt.exceptions"; +const PROFILE_LOAD_MODULE: &str = "sqlbuild.integrations.dbt._helpers.profile.load"; +const PROFILE_RENDER_MODULE: &str = "sqlbuild.integrations.dbt._helpers.profile.render"; +const PROJECT_FILE: &str = "sqlbuild_project.toml"; +const DEFAULT_OUTPUT_DIR: &str = "sqlbuild_project"; +const PROFILES_FILE: &str = "profiles.yml"; +const PROFILES_DIR_ENV_VAR: &str = "DBT_PROFILES_DIR"; +const DEBUG_SKIPPED: &str = + "dbt debug was skipped because the dbt executable was not found on PATH."; +const DEBUG_FAILED: &str = "dbt debug failed; generated SQLBuild config anyway."; + +/// The resolved dbt project folder and profiles folder of one init request. +struct InitPaths<'py> { + dbt_project_dir: Bound<'py, PyAny>, + profiles_dir: Bound<'py, PyAny>, +} + +/// Create or render a minimal SQLBuild project from a dbt project. +#[pyfunction] +#[pyo3(signature = (*, request))] +fn build_dbt_init_project<'py>(request: &Bound<'py, PyAny>) -> PyResult> { + let py: Python<'py> = request.py(); + let cwd: Bound<'py, PyAny> = request.getattr("cwd")?; + let paths: InitPaths<'py> = init_paths(request)?; + let sqb_output_dir: Bound<'py, PyAny> = request.getattr("sqb_output_dir")?; + let output_dir: Bound<'py, PyAny> = if sqb_output_dir.is_none() { + joined( + &paths.dbt_project_dir.getattr("parent")?, + DEFAULT_OUTPUT_DIR, + )? + } else { + resolved(&cwd, &sqb_output_dir)? + }; + let project_file: Bound<'py, PyAny> = joined(&output_dir, PROJECT_FILE)?; + if truthy(&project_file, "exists")? && !request.getattr("overwrite")?.is_truthy()? { + return Err(profile_error( + py, + format!( + "{} already exists. Pass --overwrite to replace it.", + project_file.str()? + ), + )); + } + if truthy(&output_dir, "exists")? && !truthy(&output_dir, "is_dir")? { + return Err(profile_error( + py, + format!( + "SQLBuild output path is not a directory: {}", + output_dir.str()? + ), + )); + } + let callbacks: Bound<'py, PyAny> = request.getattr("progress_callbacks")?; + progress(&callbacks, "start", "Inspecting dbt project and profile...")?; + let metadata: Bound<'py, PyAny> = call_with( + &imported(py, PROFILE_LOAD_MODULE, "load_dbt_project_metadata")?, + &[("project_dir", paths.dbt_project_dir.clone())], + )?; + let profile_name: Bound<'py, PyAny> = profile_name(request, &metadata)?; + let selected: Bound<'py, PyAny> = selected_output(&paths, request, &profile_name)?; + progress(&callbacks, "complete", "Inspected dbt project and profile.")?; + progress(&callbacks, "start", "Rendering dbt profile connection...")?; + let resolved_output: Bound<'py, PyAny> = call_with( + &imported( + py, + PROFILE_RENDER_MODULE, + "render_selected_dbt_profile_output", + )?, + &[ + ("selected", selected.clone()), + ("project_dir", paths.dbt_project_dir.clone()), + ("profiles_dir", paths.profiles_dir.clone()), + ], + )?; + let normalized: Bound<'py, PyAny> = call_with( + &imported(py, NATIVE_MODULE, "normalize_dbt_profile_connection")?, + &[("resolved", resolved_output)], + )?; + let target_path: Bound<'py, PyAny> = resolved( + &paths.dbt_project_dir, + &path_of(py, &metadata.getattr("target_path")?)?, + )?; + let project_name: String = metadata + .getattr("project_name")? + .extract::()? + .replace('-', "_"); + let target_name: Bound<'py, PyAny> = selected.getattr("target_name")?; + let toml: Bound<'py, PyAny> = call_with( + &imported(py, NATIVE_MODULE, "build_dbt_init_project_toml")?, + &[ + ("project_name", project_name.clone().into_bound_py_any(py)?), + ("adapter", normalized.getattr("adapter")?), + ("target_name", target_name.clone()), + ("profile_name", profile_name.clone()), + ( + "dbt_project_dir", + display_path(&paths.dbt_project_dir, &output_dir)?, + ), + ( + "profiles_dir", + display_path(&paths.profiles_dir, &output_dir)?, + ), + ("target_path", display_path(&target_path, &output_dir)?), + ("target_schema", normalized.getattr("target_schema")?), + ("target_database", normalized.getattr("target_database")?), + ], + )?; + progress(&callbacks, "complete", "Rendered dbt profile connection.")?; + let mut warnings: Vec> = normalized + .getattr("warnings")? + .try_iter()? + .collect::>()?; + if !request.getattr("skip_dbt_debug")?.is_truthy()? { + progress(&callbacks, "start", "Running dbt debug...")?; + if let Some(warning) = dbt_debug_warning(&paths, &target_name.extract::()?)? { + warnings.push(PyString::new(py, &warning).into_any()); + } + progress(&callbacks, "complete", "Finished dbt debug.")?; + } + let dry_run: Bound<'py, PyAny> = request.getattr("dry_run")?; + if !dry_run.is_truthy()? { + progress(&callbacks, "start", "Writing SQLBuild project config...")?; + let keywords: Bound<'py, PyDict> = PyDict::new(py); + keywords.set_item("parents", true)?; + keywords.set_item("exist_ok", true)?; + output_dir.call_method("mkdir", (), Some(&keywords))?; + let encoding: Bound<'py, PyDict> = PyDict::new(py); + encoding.set_item("encoding", "utf-8")?; + project_file.call_method("write_text", (&toml,), Some(&encoding))?; + progress(&callbacks, "complete", "Wrote SQLBuild project config.")?; + } + call_with( + &imported(py, DBT_MODELS, "DbtInitResult")?, + &[ + ("output_dir", output_dir), + ("project_file", project_file), + ("project_name", project_name.into_bound_py_any(py)?), + ("adapter", normalized.getattr("adapter")?), + ("target_name", target_name), + ("profile_name", profile_name), + ("toml", toml), + ("warnings", PyTuple::new(py, warnings)?.into_any()), + ("dry_run", dry_run), + ], + ) +} + +/// Validate dbt init project and profile discovery before interactive prompting. +#[pyfunction] +#[pyo3(signature = (*, request))] +fn validate_dbt_init_project(request: &Bound<'_, PyAny>) -> Result<(), PyErr> { + let paths: InitPaths<'_> = init_paths(request)?; + let metadata: Bound<'_, PyAny> = call_with( + &imported( + request.py(), + PROFILE_LOAD_MODULE, + "load_dbt_project_metadata", + )?, + &[("project_dir", paths.dbt_project_dir.clone())], + )?; + let profile_name: Bound<'_, PyAny> = profile_name(request, &metadata)?; + selected_output(&paths, request, &profile_name)?; + Ok(()) +} + +fn init_paths<'py>(request: &Bound<'py, PyAny>) -> PyResult> { + let cwd: Bound<'py, PyAny> = request.getattr("cwd")?; + let dbt_project_dir: Bound<'py, PyAny> = resolved(&cwd, &request.getattr("dbt_project_dir")?)?; + let profiles_dir: Bound<'py, PyAny> = + profiles_dir(&cwd, &dbt_project_dir, &request.getattr("profiles_dir")?)?; + Ok(InitPaths { + dbt_project_dir, + profiles_dir, + }) +} + +/// `request.profile_name or metadata.profile_name`. +fn profile_name<'py>( + request: &Bound<'py, PyAny>, + metadata: &Bound<'py, PyAny>, +) -> PyResult> { + let requested: Bound<'py, PyAny> = request.getattr("profile_name")?; + if requested.is_truthy()? { + return Ok(requested); + } + metadata.getattr("profile_name") +} + +fn selected_output<'py>( + paths: &InitPaths<'py>, + request: &Bound<'py, PyAny>, + profile_name: &Bound<'py, PyAny>, +) -> PyResult> { + let py: Python<'py> = request.py(); + let profile: Bound<'py, PyAny> = call_with( + &imported(py, PROFILE_LOAD_MODULE, "load_raw_dbt_profile")?, + &[ + ("profiles_dir", paths.profiles_dir.clone()), + ("profile_name", profile_name.clone()), + ], + )?; + call_with( + &imported(py, PROFILE_LOAD_MODULE, "select_dbt_profile_output")?, + &[ + ("profile", profile), + ("target_name", request.getattr("target_name")?), + ], + ) +} + +/// `path.expanduser()` resolved, relative paths below `root`. +fn resolved<'py>( + root: &Bound<'py, PyAny>, + path: &Bound<'py, PyAny>, +) -> PyResult> { + let expanded: Bound<'py, PyAny> = path.call_method0("expanduser")?; + if truthy(&expanded, "is_absolute")? { + return expanded.call_method0("resolve"); + } + root.call_method1("__truediv__", (expanded,))? + .call_method0("resolve") +} + +/// The explicit, `DBT_PROFILES_DIR`, project, working or default `~/.dbt` profiles folder. +fn profiles_dir<'py>( + cwd: &Bound<'py, PyAny>, + dbt_project_dir: &Bound<'py, PyAny>, + profiles_dir: &Bound<'py, PyAny>, +) -> PyResult> { + let py: Python<'py> = cwd.py(); + if !profiles_dir.is_none() { + return resolved(cwd, profiles_dir); + } + if let Some(value) = std::env::var_os(PROFILES_DIR_ENV_VAR).filter(|value| !value.is_empty()) { + return resolved(cwd, &path_of(py, &value.into_bound_py_any(py)?)?); + } + if truthy(&joined(dbt_project_dir, PROFILES_FILE)?, "exists")? { + return Ok(dbt_project_dir.clone()); + } + if truthy(&joined(cwd, PROFILES_FILE)?, "exists")? { + return cwd.call_method0("resolve"); + } + path_of(py, &PyString::new(py, "~/.dbt").into_any())?.call_method0("expanduser") +} + +/// The path as the generated config writes it: relative to `root`, else below home, else absolute. +fn display_path<'py>( + path: &Bound<'py, PyAny>, + root: &Bound<'py, PyAny>, +) -> PyResult> { + let py: Python<'py> = path.py(); + match imported(py, "os.path", "relpath")?.call1((path, root)) { + Ok(relative) => return path_of(py, &relative)?.call_method0("as_posix"), + Err(error) if is_value_error(py, &error) => {} + Err(error) => return Err(error), + } + match path.call_method1("relative_to", (root,)) { + Ok(relative) => return relative.call_method0("as_posix"), + Err(error) if is_value_error(py, &error) => {} + Err(error) => return Err(error), + } + let home: Bound<'py, PyAny> = imported(py, "pathlib", "Path")? + .call_method0("home")? + .call_method0("resolve")?; + match path.call_method1("relative_to", (home,)) { + Ok(relative) => format!( + "~/{}", + relative.call_method0("as_posix")?.extract::()? + ) + .into_bound_py_any(py), + Err(error) if is_value_error(py, &error) => path.call_method0("as_posix"), + Err(error) => Err(error), + } +} + +/// The warning `dbt debug` adds, or `None` when it passed. +fn dbt_debug_warning(paths: &InitPaths<'_>, target_name: &str) -> PyResult> { + let py: Python<'_> = paths.dbt_project_dir.py(); + let argv: Vec = vec![ + "dbt".to_owned(), + "debug".to_owned(), + "--project-dir".to_owned(), + paths.dbt_project_dir.call_method0("as_posix")?.extract()?, + "--profiles-dir".to_owned(), + paths.profiles_dir.call_method0("as_posix")?.extract()?, + "--target".to_owned(), + target_name.to_owned(), + ]; + let output: io::Result = py + .compiler_detach(|| Ok(run_dbt_process(&argv, None))) + .map_err(PyRuntimeError::new_err)?; + py.check_signals()?; + let output: DbtProcessOutput = match output { + Ok(output) => output, + Err(error) if error.kind() == io::ErrorKind::NotFound => { + return Ok(Some(DEBUG_SKIPPED.to_owned())); + } + Err(error) => return Err(launch_os_error(py, &error, (&argv, None))), + }; + if output.returncode == 0 { + return Ok(None); + } + let detail: &str = python_strip(if output.stderr.is_empty() { + &output.stdout + } else { + &output.stderr + }); + if detail.is_empty() { + return Ok(Some(DEBUG_FAILED.to_owned())); + } + Ok(Some(format!("{DEBUG_FAILED} {detail}"))) +} + +fn progress(callbacks: &Bound<'_, PyAny>, name: &str, message: &str) -> PyResult<()> { + let callback: Bound<'_, PyAny> = callbacks.getattr(name)?; + if !callback.is_none() { + callback.call1((message,))?; + } + Ok(()) +} + +fn joined<'py>(directory: &Bound<'py, PyAny>, name: &str) -> PyResult> { + directory.call_method1("__truediv__", (name,)) +} + +fn truthy(path: &Bound<'_, PyAny>, method: &str) -> PyResult { + path.call_method0(method)?.is_truthy() +} + +fn path_of<'py>(py: Python<'py>, value: &Bound<'py, PyAny>) -> PyResult> { + imported(py, "pathlib", "Path")?.call1((value,)) +} + +fn is_value_error(py: Python<'_>, error: &PyErr) -> bool { + error.is_instance_of::(py) +} + +fn profile_error(py: Python<'_>, message: String) -> PyErr { + match imported(py, DBT_EXCEPTIONS, "DbtProfileError").and_then(|class| class.call1((message,))) + { + Ok(instance) => PyErr::from_value(instance), + Err(error) => error, + } +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(build_dbt_init_project, module)?)?; + module.add_function(wrap_pyfunction!(validate_dbt_init_project, module)?) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/mod.rs index 05deced57e..24daf54eb9 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/mod.rs @@ -2,5 +2,6 @@ pub(crate) mod auto_init; pub(crate) mod command; pub(crate) mod debug; pub(crate) mod init; +pub(crate) mod init_project; pub(crate) mod registration; pub(crate) mod support; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/registration.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/registration.rs index 88cd8e8b34..50dcc0c8ee 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/registration.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/registration.rs @@ -2,9 +2,10 @@ use pyo3::prelude::{Bound, PyModule, PyResult}; -use crate::bindings::_helpers::dbt_command::{command, init}; +use crate::bindings::_helpers::dbt_command::{command, init, init_project}; pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { command::register(module)?; - init::register(module) + init::register(module)?; + init_project::register(module) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/cli.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/cli.rs index 6dfcaed30e..371202fafb 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/cli.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/cli.rs @@ -64,7 +64,7 @@ pub(crate) fn option_values(options: &Bound<'_, PyAny>) -> PyResult( +pub(crate) fn argv<'py>( py: Python<'py>, (dbt_executable, options): (&str, &Bound<'py, PyAny>), command: DbtArgvCommand, @@ -75,7 +75,7 @@ fn argv<'py>( ) } -fn sequence(values: &Bound<'_, PyAny>) -> PyResult> { +pub(crate) fn sequence(values: &Bound<'_, PyAny>) -> PyResult> { if values.is_none() { return Ok(Vec::new()); } @@ -84,7 +84,7 @@ fn sequence(values: &Bound<'_, PyAny>) -> PyResult> { /// The dbt executable, honouring a non-blank `DBT_EXECUTABLE` override. #[pyfunction] -fn resolve_dbt_executable() -> String { +pub(crate) fn resolve_dbt_executable() -> String { let value: Option = std::env::var_os(DBT_EXECUTABLE_ENV_VAR).map(|value| value.to_string_lossy().into_owned()); executable(value.as_deref()) @@ -219,6 +219,11 @@ fn ls_node<'py>( #[pyfunction] #[pyo3(signature = (*, stdout))] fn parse_dbt_ls_json_lines<'py>(py: Python<'py>, stdout: &str) -> PyResult> { + ls_nodes(py, stdout) +} + +/// The `DbtLsNode`s of `dbt ls --output json` lines. +pub(crate) fn ls_nodes<'py>(py: Python<'py>, stdout: &str) -> PyResult> { let node_class: Bound<'py, PyAny> = imported(py, NATIVE_RECORDS, "DbtLsNode")?; let mut nodes: Vec> = Vec::new(); for line in python_splitlines(stdout) { diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/event_stream.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/event_stream.rs new file mode 100644 index 0000000000..bfeedb9399 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/event_stream.rs @@ -0,0 +1,437 @@ +//! Run dbt with JSON logs and render SQLBuild-styled node rows and progress while it runs. + +use std::io::{self, BufReader, PipeReader}; +use std::path::PathBuf; +use std::process::Child; +use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender, channel}; +use std::sync::{Arc, Condvar, Mutex}; +use std::thread::JoinHandle; +use std::time::{Duration, Instant}; + +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::{ + Bound, Py, PyAny, PyAnyMethods, PyErr, PyModule, PyModuleMethods, PyResult, Python, +}; +use pyo3::sync::MutexExt; +use pyo3::types::{PyDict, PyDictMethods, PyTuple}; +use pyo3::{IntoPyObjectExt, pyfunction, wrap_pyfunction}; +use sqlbuild_core::dbt_events::main::format_dbt_active_nodes::format_dbt_active_nodes; +use sqlbuild_core::dbt_process::main::dbt_process_returncode::dbt_process_returncode; +use sqlbuild_core::dbt_process::main::read_dbt_output_line::read_dbt_output_line; +use sqlbuild_core::dbt_process::main::spawn_dbt_json_process::spawn_dbt_json_process; +use sqlbuild_core::dbt_process::main::stop_dbt_process::stop_dbt_process; +use sqlbuild_core::dbt_process::models::DbtJsonProcess; + +use crate::bindings::_helpers::build_progress::transient_status::status_reporter; +use crate::bindings::_helpers::compile_driver::python_calls::{KeywordArguments, imported, within}; +use crate::bindings::_helpers::dbt_interop::runner::{ + dbt_command_lifecycle, launch_failure, launch_os_error, report_exit, +}; +use crate::bindings::types::CompilerDetach; + +const NATIVE_MODULE: &str = "sqlbuild._native"; +const STREAM_OPTIONS: (&[&str], &[&str]) = ( + &["argv", "cwd", "stream", "use_color", "target_path"], + &[ + "display_total", + "on_node_result", + "detail_by_unique_id", + "enable_status", + ], +); +const WAITING_MESSAGE: &str = "Waiting for dbt node output..."; +/// How often the running-node progress line refreshes while dbt is silent. +const STATUS_REFRESH: Duration = Duration::from_secs(1); +/// How long cleanup waits for the progress refresher to stop. +const REFRESHER_JOIN: Duration = Duration::from_secs(1); +/// How often a blocked read rechecks for interrupt signals. +const SIGNAL_POLL: Duration = Duration::from_millis(100); + +/// The live progress line and the running nodes it lists, in first-start order. +struct StatusState { + status: Option>, + active_nodes: Vec<(String, String, f64)>, +} + +/// What the read loop and the progress refresher share. +struct Shared { + state: Mutex, + stopped: Mutex, + wake: Condvar, + epoch: Instant, +} + +/// The output stream, its colour and tty flags, and whether progress lines are drawn. +struct Output<'py> { + stream: Bound<'py, PyAny>, + use_color: bool, + enable_status: bool, +} + +/// Run dbt and render SQLBuild-styled rows from JSON events; returns the exit code and results. +#[pyfunction] +#[pyo3(signature = (**options))] +fn execute_dbt_json_event_stream<'py>( + py: Python<'py>, + options: Option>, +) -> PyResult> { + let arguments: KeywordArguments<'py> = + KeywordArguments::read(py, "execute_dbt_json_event_stream", options, STREAM_OPTIONS)?; + let enable_status: bool = + !arguments.passed("enable_status")? || arguments.flag("enable_status")?; + let output: Output<'py> = Output { + stream: arguments.value("stream")?, + use_color: arguments.flag("use_color")?, + enable_status, + }; + let events: Bound<'py, PyAny> = { + let keywords: Bound<'py, PyDict> = PyDict::new(py); + keywords.set_item("display_total", arguments.value("display_total")?)?; + keywords.set_item( + "detail_by_unique_id", + arguments.value("detail_by_unique_id")?, + )?; + keywords.set_item("use_color", output.use_color)?; + imported(py, NATIVE_MODULE, "DbtEventStream")?.call((), Some(&keywords))? + }; + let argv: Vec = arguments.value("argv")?.extract()?; + let cwd: Option = arguments.value("cwd")?.extract()?; + let target_path: Option = arguments.value("target_path")?.extract()?; + let on_node_result: Bound<'py, PyAny> = arguments.value("on_node_result")?; + let lifecycle: Bound<'py, PyAny> = dbt_command_lifecycle(py)?; + let (returncode, results) = within(&lifecycle, |lifecycle| { + let (returncode, results, process_id) = streamed( + (&argv, cwd.as_ref(), target_path.as_ref()), + &output, + (&events, &on_node_result), + )?; + report_exit(lifecycle, returncode, process_id)?; + Ok((returncode, results)) + })?; + PyTuple::new( + py, + [ + returncode.into_bound_py_any(py)?, + PyTuple::new(py, results)?.into_any(), + ], + ) +} + +/// Launch dbt, stream its lines, then always stop the refresher, the status and the process. +fn streamed<'py>( + (argv, cwd, target_path): (&[String], Option<&PathBuf>, Option<&PathBuf>), + output: &Output<'py>, + (events, on_node_result): (&Bound<'py, PyAny>, &Bound<'py, PyAny>), +) -> PyResult<(i64, Vec>, Option)> { + let py: Python<'py> = output.stream.py(); + let DbtJsonProcess { + child, + output: lines, + } = spawn_dbt_json_process( + argv, + (cwd.map(PathBuf::as_path), target_path.map(PathBuf::as_path)), + ) + .map_err(|error| { + launch_failure( + py, + launch_os_error(py, &error, (argv, cwd.map(PathBuf::as_path))), + ) + })?; + let mut child: Child = child; + let tty: bool = output.stream.call_method0("isatty")?.is_truthy()?; + let shared: Arc = Arc::new(Shared { + state: Mutex::new(StatusState { + status: None, + active_nodes: Vec::new(), + }), + stopped: Mutex::new(false), + wake: Condvar::new(), + epoch: Instant::now(), + }); + let mut refresher: Option<(JoinHandle<()>, Receiver<()>)> = None; + let mut waited: bool = false; + let result: PyResult<(i64, Vec>)> = (|| { + if output.enable_status && tty { + let status: Bound<'py, PyAny> = status_reporter(&output.stream, output.use_color)?; + status.call_method1("start", (WAITING_MESSAGE,))?; + locked(py, &shared)?.status = Some(status.unbind()); + refresher = Some(start_refresher( + &shared, + (output.stream.clone().unbind(), output.use_color), + )); + } + let mut results: Vec> = Vec::new(); + let received: Mutex>>> = Mutex::new(start_reader(lines)); + while let Some(line) = next_line(py, &received)? { + handle_line( + &line, + output, + (events, on_node_result), + (&shared, tty, &mut results), + )?; + } + let returncode: i64 = py + .compiler_detach(|| Ok(dbt_process_returncode(&mut child))) + .map_err(PyRuntimeError::new_err)? + .map_err(|error| PyErr::from(error))?; + waited = true; + py.check_signals()?; + Ok((returncode, results)) + })(); + let child: Child = cleanup(py, &shared, refresher, (child, waited)); + let (returncode, results) = result?; + Ok((returncode, results, Some(child.id()))) +} + +/// One dbt output line: progress message, node row, results and the running-node status. +fn handle_line<'py>( + line: &str, + output: &Output<'py>, + (events, on_node_result): (&Bound<'py, PyAny>, &Bound<'py, PyAny>), + (shared, tty, results): (&Shared, bool, &mut Vec>), +) -> PyResult<()> { + let py: Python<'py> = events.py(); + let update: Bound<'py, PyAny> = events.call_method1("consume", (line,))?; + if update.is_none() { + return Ok(()); + } + let message: Bound<'py, PyAny> = update.getattr("status_message")?; + if !message.is_none() && output.enable_status { + let mut state = locked(py, shared)?; + let status: Bound<'py, PyAny> = match state.status.take() { + Some(status) => { + let status: Bound<'py, PyAny> = status.into_bound(py); + status.call_method1("update", (&message,))?; + status + } + None => { + let status: Bound<'py, PyAny> = status_reporter(&output.stream, output.use_color)?; + status.call_method1("start", (&message,))?; + status + } + }; + state.status = Some(status.unbind()); + } + let result: Bound<'py, PyAny> = update.getattr("result")?; + let row: Bound<'py, PyAny> = update.getattr("row")?; + if result.is_none() || row.is_none() { + return Ok(()); + } + close_status(py, shared)?; + let finished: bool = update.getattr("finished")?.is_truthy()?; + let unique_id: String = result.getattr("unique_id")?.extract()?; + if finished { + locked(py, shared)? + .active_nodes + .retain(|(id, _, _)| *id != unique_id); + results.push(result.clone()); + if !on_node_result.is_none() { + on_node_result.call1((&result,))?; + } + } + output.stream.call_method1("write", (&row,))?; + output.stream.call_method0("flush")?; + if output.enable_status { + let mut state = locked(py, shared)?; + if !finished { + let started: f64 = shared.epoch.elapsed().as_secs_f64(); + let name: String = result.getattr("node_name")?.extract()?; + match state + .active_nodes + .iter_mut() + .find(|(id, _, _)| *id == unique_id) + { + Some(node) => *node = (unique_id, name, started), + None => state.active_nodes.push((unique_id, name, started)), + } + } + let current: Option> = state.status.take(); + state.status = refreshed_status( + py, + (&output.stream, output.use_color, tty), + shared, + (current, &state.active_nodes), + )?; + } + Ok(()) +} + +/// The redrawn running-node line; with no running node or no tty the status is dropped. +fn refreshed_status( + py: Python<'_>, + (stream, use_color, tty): (&Bound<'_, PyAny>, bool, bool), + shared: &Shared, + (current, active_nodes): (Option>, &[(String, String, f64)]), +) -> PyResult>> { + if !tty || active_nodes.is_empty() { + return Ok(None); + } + let nodes: Vec<(String, f64)> = active_nodes + .iter() + .map(|(_, name, started)| (name.clone(), *started)) + .collect(); + let message: String = format_dbt_active_nodes(&nodes, shared.epoch.elapsed().as_secs_f64()); + let status: Bound<'_, PyAny> = match current { + Some(status) => { + let status = status.into_bound(py); + status.call_method1("update", (message,))?; + status + } + None => { + let status = status_reporter(stream, use_color)?; + status.call_method1("start", (message,))?; + status + } + }; + Ok(Some(status.unbind())) +} + +fn close_status(py: Python<'_>, shared: &Shared) -> PyResult<()> { + let mut state = locked(py, shared)?; + if let Some(status) = state.status.take() { + status.bind(py).call_method0("close")?; + } + Ok(()) +} + +fn locked<'a>( + py: Python<'_>, + shared: &'a Shared, +) -> PyResult> { + shared + .state + .lock_py_attached(py) + .map_err(|_| PyRuntimeError::new_err("dbt status state is poisoned")) +} + +/// Read output lines on a thread so the waiting caller can notice interrupt signals. +fn start_reader(mut lines: BufReader) -> Receiver>> { + let (sender, receiver): ( + Sender>>, + Receiver>>, + ) = channel(); + std::thread::spawn(move || { + loop { + let line: io::Result> = read_dbt_output_line(&mut lines); + let last: bool = !matches!(line, Ok(Some(_))); + if sender.send(line).is_err() || last { + break; + } + } + }); + receiver +} + +/// The next line, raising pending interrupt signals while dbt is silent. +fn next_line( + py: Python<'_>, + received: &Mutex>>>, +) -> PyResult> { + loop { + let next: Result>, RecvTimeoutError> = py.detach(|| { + received + .lock() + .map_or(Err(RecvTimeoutError::Disconnected), |receiver| { + receiver.recv_timeout(SIGNAL_POLL) + }) + }); + match next { + Ok(line) => { + py.check_signals()?; + return line.map_err(PyErr::from); + } + Err(RecvTimeoutError::Timeout) => py.check_signals()?, + Err(RecvTimeoutError::Disconnected) => return Ok(None), + } + } +} + +/// Refresh the running-node line every second until stopped; reports when it has ended. +fn start_refresher( + shared: &Arc, + (stream, use_color): (Py, bool), +) -> (JoinHandle<()>, Receiver<()>) { + let (done, finished): (Sender<()>, Receiver<()>) = channel(); + let shared: Arc = Arc::clone(shared); + let handle: JoinHandle<()> = std::thread::spawn(move || { + while !wait_stopped(&shared) { + let refreshed: PyResult<()> = Python::attach(|py| { + let mut state = locked(py, &shared)?; + let current: Option> = state.status.take(); + state.status = refreshed_status( + py, + (stream.bind(py), use_color, true), + &shared, + (current, &state.active_nodes), + )?; + Ok(()) + }); + if refreshed.is_err() { + break; + } + } + let _sent: Result<(), std::sync::mpsc::SendError<()>> = done.send(()); + }); + (handle, finished) +} + +/// Wait one refresh period; whether cleanup asked the refresher to stop. +fn wait_stopped(shared: &Shared) -> bool { + let Ok(stopped) = shared.stopped.lock() else { + return true; + }; + if *stopped { + return true; + } + match shared.wake.wait_timeout(stopped, STATUS_REFRESH) { + Ok((stopped, _)) => *stopped, + Err(_) => true, + } +} + +/// Stop the refresher, close the status and stop dbt when it was not waited for. +fn cleanup( + py: Python<'_>, + shared: &Shared, + refresher: Option<(JoinHandle<()>, Receiver<()>)>, + (child, waited): (Child, bool), +) -> Child { + if let Ok(mut stopped) = shared.stopped.lock() { + *stopped = true; + } + shared.wake.notify_all(); + if let Some((handle, finished)) = refresher { + let finished: Mutex> = Mutex::new(finished); + let ended: bool = py.detach(|| match finished.lock() { + Ok(receiver) => match receiver.recv_timeout(REFRESHER_JOIN) { + Ok(()) => true, + Err(RecvTimeoutError::Timeout | RecvTimeoutError::Disconnected) => false, + }, + Err(_) => false, + }); + if ended { + let _joined: std::thread::Result<()> = handle.join(); + } + } + let _closed: PyResult<()> = close_status(py, shared); + if waited { + return child; + } + py.detach(|| stop_dbt_process(child, terminate)) +} + +/// Ask a process to stop as `Popen.terminate` does. +fn terminate(process_id: u32) { + let _requested: PyResult<()> = Python::attach(|py| { + if cfg!(windows) { + return Ok(()); + } + let signal: Bound<'_, PyAny> = imported(py, "signal", "SIGTERM")?; + imported(py, "os", "kill")?.call1((process_id, signal))?; + Ok(()) + }); +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(execute_dbt_json_event_stream, module)?) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/mod.rs index f60b21c5ef..4101df8eb8 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/mod.rs @@ -1,5 +1,6 @@ pub(crate) mod cli; pub(crate) mod constants; +pub(crate) mod event_stream; pub(crate) mod execution; pub(crate) mod graph; pub(crate) mod manifest; @@ -13,6 +14,7 @@ pub(crate) mod profile; pub(crate) mod profile_connection; pub(crate) mod reference_resolver; pub(crate) mod registration; +pub(crate) mod runner; pub(crate) mod runtime_config; pub(crate) mod runtime_execution; pub(crate) mod runtime_results; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/pipeline.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/pipeline.rs index 9fa36ef698..83de2c34b8 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/pipeline.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/pipeline.rs @@ -358,11 +358,7 @@ fn debug_dbt_from_project(py: Python<'_>, options: Option>) -> let mut runner = args.value("dbt_runner")?; if !runner.is_truthy()? { runner = call_with( - &imported( - py, - "sqlbuild.integrations.dbt.classes.dbt_runner", - "DbtRunner", - )?, + &imported(py, "sqlbuild._native", "DbtRunner")?, &[("dbt_executable", args.value("dbt_executable")?)], )?; } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/pipeline_prologue.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/pipeline_prologue.rs index b12d63d053..939f0d94b2 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/pipeline_prologue.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/pipeline_prologue.rs @@ -97,11 +97,7 @@ pub(crate) fn resolve_dbt_execution_invocation<'py>( let mut runner = request.getattr("dbt_runner")?; if !runner.is_truthy()? { runner = call_with( - &imported( - py, - "sqlbuild.integrations.dbt.classes.dbt_runner", - "DbtRunner", - )?, + &imported(py, "sqlbuild._native", "DbtRunner")?, &[("dbt_executable", executable.clone())], )?; } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/registration.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/registration.rs index 0b9b4076f2..02bb776394 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/registration.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/registration.rs @@ -22,5 +22,7 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { super::plan_output::register(module)?; super::sqlbuild_work::register(module)?; super::pipeline_phases::register(module)?; + super::runner::register(module)?; + super::event_stream::register(module)?; super::pipeline::register(module) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/runner.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/runner.rs new file mode 100644 index 0000000000..5f66035ae2 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/runner.rs @@ -0,0 +1,296 @@ +//! The dbt CLI runner: captured dbt commands with operation lifecycle events and `dbt ls` memo. + +use std::collections::HashMap; +use std::io; +use std::path::PathBuf; +use std::sync::Mutex; + +use pyo3::exceptions::{PyOSError, PyRuntimeError}; +use pyo3::prelude::{ + Bound, Py, PyAny, PyAnyMethods, PyErr, PyModule, PyModuleMethods, PyResult, Python, +}; +use pyo3::sync::MutexExt; +use pyo3::types::{PyDict, PyDictMethods, PyTuple}; +use pyo3::{pyclass, pymethods}; +use sqlbuild_core::dbt_cli::models::DbtArgvCommand; +use sqlbuild_core::dbt_process::main::run_dbt_process::run_dbt_process; +use sqlbuild_core::dbt_process::models::DbtProcessOutput; + +use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported, within}; +use crate::bindings::_helpers::dbt_interop::cli::{ + argv, ls_nodes, resolve_dbt_executable, sequence, +}; +use crate::bindings::types::CompilerDetach; + +const DBT_MODELS: &str = "sqlbuild.integrations.dbt.models"; +const DBT_EXCEPTIONS: &str = "sqlbuild.integrations.dbt.exceptions"; +const LIFECYCLE_MODULE: &str = "sqlbuild.runtime.observability.classes.operation_lifecycle"; +const LAUNCH_FAILURE: &str = "failed to execute dbt"; +const POISONED_CACHE: &str = "dbt ls cache is poisoned"; + +/// Small dbt CLI runner with per-invocation `dbt ls` memoization. +#[pyclass(module = "sqlbuild._native", frozen, subclass)] +pub(crate) struct DbtRunner { + #[pyo3(get)] + dbt_executable: String, + #[pyo3(get)] + invoker: Option>, + ls_cache: Mutex, Py>>, +} + +#[pymethods] +impl DbtRunner { + #[new] + #[pyo3(signature = (*, dbt_executable = None, invoker = None))] + fn new(dbt_executable: Option, invoker: Option>) -> Self { + Self { + dbt_executable: dbt_executable + .filter(|value| !value.is_empty()) + .unwrap_or_else(resolve_dbt_executable), + invoker, + ls_cache: Mutex::new(HashMap::new()), + } + } + + /// Run dbt compile. + #[pyo3(signature = (*, options, full_refresh = false))] + fn compile<'py>( + &self, + options: &Bound<'py, PyAny>, + full_refresh: bool, + ) -> PyResult> { + let command: Bound<'py, PyTuple> = argv( + options.py(), + (&self.dbt_executable, options), + DbtArgvCommand::Compile { full_refresh }, + )?; + self.run(command, &options.getattr("project_dir")?) + } + + /// Run dbt deps. + #[pyo3(signature = (*, options))] + fn deps<'py>(&self, options: &Bound<'py, PyAny>) -> PyResult> { + let command: Bound<'py, PyTuple> = argv( + options.py(), + (&self.dbt_executable, options), + DbtArgvCommand::Deps, + )?; + self.run(command, &options.getattr("project_dir")?) + } + + /// Run dbt debug. + #[pyo3(signature = (*, options, args = None))] + fn debug<'py>( + &self, + options: &Bound<'py, PyAny>, + args: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + let args: Vec = args.map_or_else(|| Ok(Vec::new()), sequence)?; + let command: Bound<'py, PyTuple> = argv( + options.py(), + (&self.dbt_executable, options), + DbtArgvCommand::Debug { args }, + )?; + self.run(command, &options.getattr("project_dir")?) + } + + /// Run dbt ls and parse selected nodes, reusing an earlier identical listing. + #[pyo3(signature = (*, options, select = None, exclude = None, resource_types = None))] + fn ls<'py>( + &self, + options: &Bound<'py, PyAny>, + select: Option<&Bound<'py, PyAny>>, + exclude: Option<&Bound<'py, PyAny>>, + resource_types: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + let py: Python<'py> = options.py(); + let listed = |values: Option<&Bound<'py, PyAny>>| -> PyResult> { + values.map_or_else(|| Ok(Vec::new()), sequence) + }; + let command: Bound<'py, PyTuple> = argv( + py, + (&self.dbt_executable, options), + DbtArgvCommand::Ls { + select: listed(select)?, + exclude: listed(exclude)?, + resource_types: listed(resource_types)?, + }, + )?; + let key: Vec = command.extract()?; + if let Some(cached) = self.cached_listing(py, &key)? { + return Ok(cached); + } + let result: Bound<'py, PyAny> = self.run(command, &options.getattr("project_dir")?)?; + let stdout: String = result.getattr("stdout")?.extract()?; + let listing: Bound<'py, PyAny> = call_with( + &imported(py, DBT_MODELS, "DbtLsResult")?, + &[ + ("nodes", ls_nodes(py, &stdout)?.into_any()), + ("command", result), + ], + )?; + self.ls_cache + .lock_py_attached(py) + .map_err(|_| PyRuntimeError::new_err(POISONED_CACHE))? + .insert(key, listing.clone().unbind()); + Ok(listing) + } + + /// Run an explicit dbt argv. + #[pyo3(signature = (*, argv, cwd))] + fn invoke<'py>( + &self, + argv: &Bound<'py, PyAny>, + cwd: &Bound<'py, PyAny>, + ) -> PyResult> { + let command: Bound<'py, PyTuple> = + PyTuple::new(argv.py(), argv.try_iter()?.collect::>>()?)?; + self.run(command, cwd) + } +} + +impl DbtRunner { + fn cached_listing<'py>( + &self, + py: Python<'py>, + key: &[String], + ) -> PyResult>> { + let cache = self + .ls_cache + .lock_py_attached(py) + .map_err(|_| PyRuntimeError::new_err(POISONED_CACHE))?; + Ok(cache + .get(key) + .map(|value| value.clone_ref(py).into_bound(py))) + } + + fn run<'py>( + &self, + command: Bound<'py, PyTuple>, + cwd: &Bound<'py, PyAny>, + ) -> PyResult> { + let py: Python<'py> = command.py(); + if let Some(invoker) = &self.invoker { + return call_with( + invoker.bind(py), + &[("argv", command.into_any()), ("cwd", cwd.clone())], + ); + } + let arguments: Vec = command.extract()?; + let directory: Option = cwd.extract()?; + let lifecycle: Bound<'py, PyAny> = dbt_command_lifecycle(py)?; + let output: PyResult = within(&lifecycle, |lifecycle| { + let output: io::Result = py + .compiler_detach(|| Ok(run_dbt_process(&arguments, directory.as_deref()))) + .map_err(PyRuntimeError::new_err)?; + py.check_signals()?; + let output: DbtProcessOutput = output + .map_err(|error| launch_os_error(py, &error, (&arguments, directory.as_deref())))?; + report_exit(lifecycle, output.returncode, None)?; + Ok(output) + }); + let output: DbtProcessOutput = output.map_err(|error| launch_failure(py, error))?; + call_with( + &imported(py, DBT_MODELS, "DbtCommandResult")?, + &[ + ("argv", command.into_any()), + ("returncode", output.returncode.into_pyobject_any(py)?), + ("stdout", output.stdout.into_pyobject_any(py)?), + ("stderr", output.stderr.into_pyobject_any(py)?), + ], + ) + } +} + +/// Shorthand for converting owned Rust values into Python values. +trait IntoAny { + fn into_pyobject_any<'py>(self, py: Python<'py>) -> PyResult>; +} + +impl pyo3::IntoPyObject<'py>> IntoAny for T { + fn into_pyobject_any<'py>(self, py: Python<'py>) -> PyResult> { + pyo3::IntoPyObjectExt::into_bound_py_any(self, py) + } +} + +/// `OperationLifecycle(operation_kind="subprocess", operation_name="dbt_command", ...)`. +pub(crate) fn dbt_command_lifecycle(py: Python<'_>) -> PyResult> { + let options: Bound<'_, PyDict> = PyDict::new(py); + options.set_item("operation_kind", "subprocess")?; + options.set_item("operation_name", "dbt_command")?; + options.set_item("auto_fail_base_exceptions", false)?; + imported(py, LIFECYCLE_MODULE, "OperationLifecycle")?.call((), Some(&options)) +} + +/// Complete a zero exit, or fail with the exit code and any signal number. +pub(crate) fn report_exit( + lifecycle: &Bound<'_, PyAny>, + returncode: i64, + process_id: Option, +) -> PyResult<()> { + let py: Python<'_> = lifecycle.py(); + let options: Bound<'_, PyDict> = PyDict::new(py); + if let Some(process_id) = process_id { + options.set_item("process_id", process_id)?; + } + if returncode == 0 { + options.set_item("exit_code", 0)?; + lifecycle.call_method("completed", (), Some(&options))?; + return Ok(()); + } + options.set_item("error_code", format!("exit_{returncode}"))?; + options.set_item("exit_code", returncode)?; + options.set_item( + "signal_number", + if returncode < 0 { + Some(-returncode) + } else { + None + }, + )?; + lifecycle.call_method("failed", (), Some(&options))?; + Ok(()) +} + +/// The `OSError` subclass Python's `subprocess` raises for a failed launch. +pub(crate) fn launch_os_error( + py: Python<'_>, + error: &io::Error, + (arguments, cwd): (&[String], Option<&std::path::Path>), +) -> PyErr { + let Some(code) = error.raw_os_error() else { + return PyOSError::new_err(error.to_string()); + }; + let filename: String = match cwd { + Some(directory) if !directory.is_dir() => directory.display().to_string(), + _ => arguments.first().cloned().unwrap_or_default(), + }; + let message: PyResult = imported(py, "os", "strerror") + .and_then(|strerror| strerror.call1((code,))) + .and_then(|text| text.extract()); + match message { + Ok(message) => PyOSError::new_err((code, message, filename)), + Err(failure) => failure, + } +} + +/// `DbtInteropRuntimeError("failed to execute dbt", help=str(error))` from a launch `OSError`. +pub(crate) fn launch_failure(py: Python<'_>, error: PyErr) -> PyErr { + if !error.is_instance_of::(py) { + return error; + } + let created = || -> PyResult { + let options: Bound<'_, PyDict> = PyDict::new(py); + options.set_item("help", error.value(py).str()?)?; + let instance: Bound<'_, PyAny> = imported(py, DBT_EXCEPTIONS, "DbtInteropRuntimeError")? + .call((LAUNCH_FAILURE,), Some(&options))?; + let raised: PyErr = PyErr::from_value(instance); + raised.set_cause(py, Some(error.clone_ref(py))); + Ok(raised) + }; + created().unwrap_or_else(|failure| failure) +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_class::() +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/runtime_execution.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/runtime_execution.rs index 9fc861a255..bf09afd152 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/runtime_execution.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/runtime_execution.rs @@ -174,11 +174,7 @@ fn execute_dbt_commands<'py>( target = project_dir.call_method1("__truediv__", ("target",))?; } let (returncode, streamed): (Bound<'py, PyAny>, Bound<'py, PyAny>) = call_with( - &imported( - py, - "sqlbuild.integrations.dbt.main.runtime._execute_json_event_stream", - "execute_dbt_json_event_stream", - )?, + &imported(py, "sqlbuild._native", "execute_dbt_json_event_stream")?, &[ ("argv", argv), ("cwd", project_dir), diff --git a/fensu.toml b/fensu.toml index 8476ce1a4b..86eca090dd 100644 --- a/fensu.toml +++ b/fensu.toml @@ -43,6 +43,11 @@ rule = "FFR309" path = "src/sqlbuild/spec/contracts/__init__.py" reason = "Native bindings own configuration resolution; this package only declares the configuration models, constants and errors callers import." +[[targets.python.rule_exceptions]] +rule = "FFR309" +path = "src/sqlbuild/integrations/dbt/__init__.py" +reason = "Native bindings own the dbt interop commands, the dbt runner, its JSON event stream and project initialisation; this package keeps the dbt models, types and errors and the PyYAML and Jinja profile glue they call." + [[targets.python.rule_exceptions]] rule = "FFR309" path = "src/sqlbuild/compiler/frontier/__init__.py" diff --git a/python-boundary.toml b/python-boundary.toml index 99d850d278..c8f7239217 100644 --- a/python-boundary.toml +++ b/python-boundary.toml @@ -2,7 +2,7 @@ # migration work that moves to Rust. `remaining_lines` must equal the current count of that work: # a PR that moves code lowers it, and adding movable Python fails the check. # Native failures must surface as errors; `handled_native_errors` lists the reviewed exceptions. -remaining_lines = 9601 +remaining_lines = 9599 [[stays]] path = "executor" @@ -294,21 +294,6 @@ path = "integrations/ingestr" reason = "user-code" note = "Loader runtime that runs inside the executor load phase." -[[stays]] -path = "integrations/dbt/_helpers/runtime/event_stream.py" -reason = "user-code" -note = "Runs the dbt subprocess over the user's dbt project and drives its progress threads." - -[[stays]] -path = "integrations/dbt/main/runtime/_execute_json_event_stream.py" -reason = "user-code" -note = "Entry point of the dbt subprocess runner." - -[[stays]] -path = "integrations/dbt/classes/dbt_runner.py" -reason = "user-code" -note = "Runs dbt compile, deps, debug and ls as subprocesses over the user's dbt project." - [[stays]] path = "integrations/dbt/_helpers/profile/load.py" reason = "user-code" @@ -334,11 +319,6 @@ path = "cli/commands/_helpers/entry/compute_log_diagnostics.py" reason = "user-code" note = "Logs compute log capture failures through Python logging during capture." -[[stays]] -path = "integrations/dbt/_helpers/profile/init.py" -reason = "user-code" -note = "Initializes a SQLBuild project from the user's dbt project and runs dbt debug as a subprocess." - [[stays]] path = "cli/compile_reuse/_helpers/provider_settings.py" reason = "user-code" diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index d118a64e3f..2902ab48ef 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -307,14 +307,29 @@ from sqlbuild.executor.python_nodes.models import ( from sqlbuild.executor.run.models import ModelExecutionResult from sqlbuild.executor.scenario.models import ScenarioRunResult, ScenarioSnapshotCaptureRunResult from sqlbuild.executor.testing.models import SqlTestExecutionResult, StepResult -from sqlbuild.integrations.dbt.classes.dbt_runner import DbtRunner from sqlbuild.integrations.dbt.models import ( DbtCliOptions, DbtCommandResult, + DbtInitRequest, + DbtInitResult, + DbtInteropCommandArgs, + DbtInteropCompiledProject, + DbtInteropExecutionRequest, + DbtInteropInvocation, + DbtInteropParsedArgs, + DbtInteropPlan, + DbtInteropPlanResolution, + DbtInteropRoutedArgs, + DbtInteropSelectionResult, + DbtLsNode, + DbtLsResult, + DbtManifestIndex, + DbtManifestModel, DbtNodeExecutionResult, NormalizedDbtProfileConnection, ResolvedDbtProfileOutput, ) +from sqlbuild.integrations.dbt.types import DbtInteropCommand, DbtInvoker from sqlbuild.lint.models import ( CompileFacts, ExpandedLintFile, @@ -3075,6 +3090,42 @@ def plan_dbt_interop_command( command_args: DbtInteropCommandArgs, ) -> DbtInteropPlan: ... def resolve_dbt_executable() -> str: ... + +class DbtRunner: + dbt_executable: str + invoker: DbtInvoker | None + def __init__( + self, *, dbt_executable: str | None = None, invoker: DbtInvoker | None = None + ) -> None: ... + def compile( + self, *, options: DbtCliOptions, full_refresh: bool = False + ) -> DbtCommandResult: ... + def deps(self, *, options: DbtCliOptions) -> DbtCommandResult: ... + def debug(self, *, options: DbtCliOptions, args: Sequence[str] = ()) -> DbtCommandResult: ... + def ls( + self, + *, + options: DbtCliOptions, + select: Sequence[str] = (), + exclude: Sequence[str] = (), + resource_types: Sequence[str] = (), + ) -> DbtLsResult: ... + def invoke(self, *, argv: tuple[str, ...], cwd: Path | None) -> DbtCommandResult: ... + +def execute_dbt_json_event_stream( + *, + argv: tuple[str, ...], + cwd: Path | None, + stream: TextIO, + use_color: bool, + target_path: Path | None, + display_total: int | None = None, + on_node_result: Callable[[DbtNodeExecutionResult], None] | None = None, + detail_by_unique_id: dict[str, str] | None = None, + enable_status: bool = True, +) -> tuple[int, tuple[DbtNodeExecutionResult, ...]]: ... +def build_dbt_init_project(*, request: DbtInitRequest) -> DbtInitResult: ... +def validate_dbt_init_project(*, request: DbtInitRequest) -> None: ... def build_dbt_compile_argv( *, dbt_executable: str, options: DbtCliOptions, full_refresh: bool = False ) -> tuple[str, ...]: ... diff --git a/src/sqlbuild/integrations/dbt/_helpers/profile/init.py b/src/sqlbuild/integrations/dbt/_helpers/profile/init.py deleted file mode 100644 index 7eb32dcdf6..0000000000 --- a/src/sqlbuild/integrations/dbt/_helpers/profile/init.py +++ /dev/null @@ -1,207 +0,0 @@ -"""Initialize SQLBuild projects from dbt projects.""" - -from __future__ import annotations - -import os -import subprocess -from pathlib import Path - -from sqlbuild._native import build_dbt_init_project_toml, normalize_dbt_profile_connection -from sqlbuild.integrations.dbt._helpers.profile.load import ( - default_profiles_dir, - load_dbt_project_metadata, - load_raw_dbt_profile, - select_dbt_profile_output, -) -from sqlbuild.integrations.dbt._helpers.profile.render import render_selected_dbt_profile_output -from sqlbuild.integrations.dbt.exceptions import DbtProfileError -from sqlbuild.integrations.dbt.models import ( - DbtInitProgressCallbacks, - DbtInitRequest, - DbtInitResult, - DbtProjectProfileMetadata, - NormalizedDbtProfileConnection, - ResolvedDbtProfileOutput, - SelectedDbtProfileOutput, -) - - -def build_dbt_init_project(*, request: DbtInitRequest) -> DbtInitResult: - """Create or render a minimal SQLBuild project from a dbt project.""" - - dbt_project_dir: Path = _resolve_path(root=request.cwd, path=request.dbt_project_dir) - profiles_dir: Path = _resolve_profiles_dir( - cwd=request.cwd, - dbt_project_dir=dbt_project_dir, - profiles_dir=request.profiles_dir, - ) - output_dir: Path = ( - _resolve_path(root=request.cwd, path=request.sqb_output_dir) - if request.sqb_output_dir is not None - else dbt_project_dir.parent / "sqlbuild_project" - ) - project_file: Path = output_dir / "sqlbuild_project.toml" - if project_file.exists() and not request.overwrite: - raise DbtProfileError(f"{project_file} already exists. Pass --overwrite to replace it.") - if output_dir.exists() and not output_dir.is_dir(): - raise DbtProfileError(f"SQLBuild output path is not a directory: {output_dir}") - - _start_progress( - callbacks=request.progress_callbacks, message="Inspecting dbt project and profile..." - ) - metadata: DbtProjectProfileMetadata = load_dbt_project_metadata(project_dir=dbt_project_dir) - profile_name: str = request.profile_name or metadata.profile_name - selected: SelectedDbtProfileOutput = select_dbt_profile_output( - profile=load_raw_dbt_profile(profiles_dir=profiles_dir, profile_name=profile_name), - target_name=request.target_name, - ) - _complete_progress( - callbacks=request.progress_callbacks, message="Inspected dbt project and profile." - ) - - _start_progress( - callbacks=request.progress_callbacks, message="Rendering dbt profile connection..." - ) - resolved: ResolvedDbtProfileOutput = render_selected_dbt_profile_output( - selected=selected, - project_dir=dbt_project_dir, - profiles_dir=profiles_dir, - ) - normalized: NormalizedDbtProfileConnection = normalize_dbt_profile_connection(resolved=resolved) - target_path: Path = _resolve_path(root=dbt_project_dir, path=Path(metadata.target_path)) - toml: str = build_dbt_init_project_toml( - project_name=metadata.project_name.replace("-", "_"), - adapter=normalized.adapter, - target_name=selected.target_name, - profile_name=profile_name, - dbt_project_dir=_display_path(path=dbt_project_dir, root=output_dir), - profiles_dir=_display_path(path=profiles_dir, root=output_dir), - target_path=_display_path(path=target_path, root=output_dir), - target_schema=normalized.target_schema, - target_database=normalized.target_database, - ) - _complete_progress( - callbacks=request.progress_callbacks, message="Rendered dbt profile connection." - ) - warnings: list[str] = list(normalized.warnings) - if not request.skip_dbt_debug: - _start_progress(callbacks=request.progress_callbacks, message="Running dbt debug...") - warnings.extend( - _dbt_debug_warnings( - dbt_project_dir=dbt_project_dir, - profiles_dir=profiles_dir, - target_name=selected.target_name, - ) - ) - _complete_progress(callbacks=request.progress_callbacks, message="Finished dbt debug.") - if not request.dry_run: - _start_progress( - callbacks=request.progress_callbacks, message="Writing SQLBuild project config..." - ) - output_dir.mkdir(parents=True, exist_ok=True) - project_file.write_text(toml, encoding="utf-8") - _complete_progress( - callbacks=request.progress_callbacks, message="Wrote SQLBuild project config." - ) - return DbtInitResult( - output_dir=output_dir, - project_file=project_file, - project_name=metadata.project_name.replace("-", "_"), - adapter=normalized.adapter, - target_name=selected.target_name, - profile_name=profile_name, - toml=toml, - warnings=tuple(warnings), - dry_run=request.dry_run, - ) - - -def validate_dbt_init_project(*, request: DbtInitRequest) -> None: - """Validate dbt init project/profile discovery before interactive prompting.""" - - dbt_project_dir: Path = _resolve_path(root=request.cwd, path=request.dbt_project_dir) - profiles_dir: Path = _resolve_profiles_dir( - cwd=request.cwd, - dbt_project_dir=dbt_project_dir, - profiles_dir=request.profiles_dir, - ) - metadata: DbtProjectProfileMetadata = load_dbt_project_metadata(project_dir=dbt_project_dir) - profile_name: str = request.profile_name or metadata.profile_name - select_dbt_profile_output( - profile=load_raw_dbt_profile(profiles_dir=profiles_dir, profile_name=profile_name), - target_name=request.target_name, - ) - - -def _resolve_path(*, root: Path, path: Path) -> Path: - expanded: Path = path.expanduser() - if expanded.is_absolute(): - return expanded.resolve() - return (root / expanded).resolve() - - -def _resolve_profiles_dir(*, cwd: Path, dbt_project_dir: Path, profiles_dir: Path | None) -> Path: - if profiles_dir is not None: - return _resolve_path(root=cwd, path=profiles_dir) - env_profiles_dir: str | None = os.environ.get("DBT_PROFILES_DIR") - if env_profiles_dir: - return _resolve_path(root=cwd, path=Path(env_profiles_dir)) - if (dbt_project_dir / "profiles.yml").exists(): - return dbt_project_dir - if (cwd / "profiles.yml").exists(): - return cwd.resolve() - return default_profiles_dir() - - -def _start_progress(*, callbacks: DbtInitProgressCallbacks, message: str) -> None: - if callbacks.start is not None: - callbacks.start(message) - - -def _complete_progress(*, callbacks: DbtInitProgressCallbacks, message: str) -> None: - if callbacks.complete is not None: - callbacks.complete(message) - - -def _display_path(*, path: Path, root: Path) -> str: - try: - return Path(os.path.relpath(path, root)).as_posix() - except ValueError: - pass - try: - return path.relative_to(root).as_posix() - except ValueError: - home: Path = Path.home().resolve() - try: - return "~/" + path.relative_to(home).as_posix() - except ValueError: - return path.as_posix() - - -def _dbt_debug_warnings( - *, dbt_project_dir: Path, profiles_dir: Path, target_name: str -) -> tuple[str, ...]: - try: - result: subprocess.CompletedProcess[str] = subprocess.run( - ( - "dbt", - "debug", - "--project-dir", - dbt_project_dir.as_posix(), - "--profiles-dir", - profiles_dir.as_posix(), - "--target", - target_name, - ), - capture_output=True, - check=False, - text=True, - ) - except FileNotFoundError: - return ("dbt debug was skipped because the dbt executable was not found on PATH.",) - if result.returncode == 0: - return () - detail: str = (result.stderr or result.stdout).strip() - if detail: - return ("dbt debug failed; generated SQLBuild config anyway. " + detail,) - return ("dbt debug failed; generated SQLBuild config anyway.",) diff --git a/src/sqlbuild/integrations/dbt/_helpers/runtime/event_stream.py b/src/sqlbuild/integrations/dbt/_helpers/runtime/event_stream.py deleted file mode 100644 index 754fc6c7ca..0000000000 --- a/src/sqlbuild/integrations/dbt/_helpers/runtime/event_stream.py +++ /dev/null @@ -1,405 +0,0 @@ -"""dbt JSON event parsing and streaming helpers.""" - -from __future__ import annotations - -import os -import subprocess -import threading -import time -from collections.abc import Callable -from dataclasses import dataclass, field -from pathlib import Path -from typing import TextIO - -from sqlbuild._native import ( - DbtEventStream, - DbtEventStreamUpdate, - TransientStatusReporter, - format_dbt_active_node_status, -) -from sqlbuild.integrations.dbt.exceptions import DbtInteropRuntimeError -from sqlbuild.integrations.dbt.models import DbtNodeExecutionResult -from sqlbuild.runtime.observability.classes.operation_lifecycle import OperationLifecycle - -_DBT_STATUS_REFRESH_SECONDS: float = 1.0 - - -@dataclass -class _DbtStreamState: - events: DbtEventStream - results: list[DbtNodeExecutionResult] = field(default_factory=list) - active_nodes: dict[str, tuple[str, float]] = field(default_factory=dict) - use_color: bool = False - - -def execute_dbt_json_event_stream( - *, - argv: tuple[str, ...], - cwd: Path | None, - stream: TextIO, - use_color: bool, - target_path: Path | None, - display_total: int | None = None, - on_node_result: Callable[[DbtNodeExecutionResult], None] | None = None, - detail_by_unique_id: dict[str, str] | None = None, - enable_status: bool = True, -) -> tuple[int, tuple[DbtNodeExecutionResult, ...]]: - """Run dbt and render SQLBuild-styled rows from JSON events.""" - - with OperationLifecycle( - operation_kind="subprocess", - operation_name="dbt_command", - auto_fail_base_exceptions=False, - ) as lifecycle: - returncode, results, process_id = _execute_dbt_json_event_stream( - argv=argv, - cwd=cwd, - stream=stream, - use_color=use_color, - target_path=target_path, - display_total=display_total, - on_node_result=on_node_result, - detail_by_unique_id=detail_by_unique_id, - enable_status=enable_status, - ) - if returncode == 0: - lifecycle.completed(exit_code=0, process_id=process_id) - else: - lifecycle.failed( - error_code=f"exit_{returncode}", - exit_code=returncode, - process_id=process_id, - signal_number=-returncode if returncode < 0 else None, - ) - return returncode, results - - -def _execute_dbt_json_event_stream( - *, - argv: tuple[str, ...], - cwd: Path | None, - stream: TextIO, - use_color: bool, - target_path: Path | None, - display_total: int | None, - on_node_result: Callable[[DbtNodeExecutionResult], None] | None, - detail_by_unique_id: dict[str, str] | None, - enable_status: bool, -) -> tuple[int, tuple[DbtNodeExecutionResult, ...], int | None]: - state: _DbtStreamState = _DbtStreamState( - events=DbtEventStream( - display_total=display_total, - detail_by_unique_id=detail_by_unique_id, - use_color=use_color, - ), - use_color=use_color, - ) - status_box: dict[str, TransientStatusReporter | None] = {"status": None} - status_lock: threading.Lock = threading.Lock() - status_stop: threading.Event = threading.Event() - status_thread: threading.Thread | None = None - process: subprocess.Popen[str] | None = None - process_waited = False - try: - process = _launch_dbt_process( - argv=argv, - cwd=cwd, - target_path=target_path, - ) - status_box["status"] = ( - _start_dbt_status(stream=stream, use_color=use_color) if enable_status else None - ) - status_thread = _start_active_node_status_refresher( - stream=stream, - use_color=use_color, - enabled=enable_status, - active_nodes=state.active_nodes, - status_box=status_box, - status_lock=status_lock, - status_stop=status_stop, - ) - if process.stdout is not None: - try: - for line in process.stdout: - state, status_box = _handle_dbt_output_line( - line=line, - stream=stream, - enable_status=enable_status, - state=state, - status_box=status_box, - status_lock=status_lock, - on_node_result=on_node_result, - ) - finally: - try: - process.stdout.close() - except BaseException: - pass - returncode: int = process.wait() - process_waited = True - process_id: object = getattr(process, "pid", None) - return ( - returncode, - tuple(state.results), - process_id - if isinstance(process_id, int) and not isinstance(process_id, bool) and process_id >= 0 - else None, - ) - finally: - _cleanup_dbt_runtime( - process=process, - process_waited=process_waited, - status_stop=status_stop, - status_thread=status_thread, - status_box=status_box, - status_lock=status_lock, - ) - - -def _handle_dbt_output_line( - *, - line: str, - stream: TextIO, - enable_status: bool, - state: _DbtStreamState, - status_box: dict[str, TransientStatusReporter | None], - status_lock: threading.Lock, - on_node_result: Callable[[DbtNodeExecutionResult], None] | None, -) -> tuple[_DbtStreamState, dict[str, TransientStatusReporter | None]]: - update: DbtEventStreamUpdate | None = state.events.consume(line) - if update is None: - return state, status_box - if update.status_message is not None and enable_status: - with status_lock: - status_box["status"] = _update_message_status( - status=status_box["status"], - message=update.status_message, - stream=stream, - use_color=state.use_color, - ) - result: DbtNodeExecutionResult | None = update.result - if result is None or update.row is None: - return state, status_box - status_box = _close_status(status_box=status_box, status_lock=status_lock) - if update.finished: - with status_lock: - state.active_nodes.pop(result.unique_id, None) - state.results.append(result) - if on_node_result is not None: - on_node_result(result) - stream.write(update.row) - stream.flush() - if enable_status: - with status_lock: - if not update.finished: - state.active_nodes[result.unique_id] = (result.node_name, time.monotonic()) - status_box["status"] = _update_active_node_status( - status=status_box["status"], - active_nodes=state.active_nodes, - stream=stream, - use_color=state.use_color, - ) - return state, status_box - - -def _cleanup_dbt_runtime( - *, - process: subprocess.Popen[str] | None, - process_waited: bool, - status_stop: threading.Event, - status_thread: threading.Thread | None, - status_box: dict[str, TransientStatusReporter | None], - status_lock: threading.Lock, -) -> None: - try: - status_stop.set() - except BaseException: - pass - if status_thread is not None: - try: - status_thread.join(timeout=1) - except BaseException: - pass - try: - _ = _close_status(status_box=status_box, status_lock=status_lock) - except BaseException: - pass - if process is not None and not process_waited: - _stop_dbt_process(process=process) - - -def _stop_dbt_process(*, process: subprocess.Popen[str]) -> None: - returncode: int | None = None - try: - returncode = process.poll() - except BaseException: - pass - if returncode is None: - try: - process.terminate() - except BaseException: - pass - try: - _ = process.wait(timeout=1) - return - except subprocess.TimeoutExpired: - try: - process.kill() - except BaseException: - pass - except BaseException: - try: - process.kill() - except BaseException: - pass - try: - _ = process.wait(timeout=1) - except BaseException: - pass - - -def _launch_dbt_process( - *, - argv: tuple[str, ...], - cwd: Path | None, - target_path: Path | None, -) -> subprocess.Popen[str]: - try: - return subprocess.Popen( - argv, - cwd=cwd, - env=_build_dbt_json_env(target_path=target_path), - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - bufsize=1, - ) - except OSError as error: - raise DbtInteropRuntimeError( - "failed to execute dbt", - help=str(error), - ) from error - - -def _start_dbt_status(*, stream: TextIO, use_color: bool) -> TransientStatusReporter | None: - if not stream.isatty(): - return None - status: TransientStatusReporter = TransientStatusReporter( - stream=stream, - use_color=use_color, - ) - status.start("Waiting for dbt node output...") - return status - - -def _start_active_node_status_refresher( - *, - stream: TextIO, - use_color: bool, - enabled: bool, - active_nodes: dict[str, tuple[str, float]], - status_box: dict[str, TransientStatusReporter | None], - status_lock: threading.Lock, - status_stop: threading.Event, - refresh_seconds: float | None = None, -) -> threading.Thread | None: - if not enabled or not stream.isatty(): - return None - actual_refresh_seconds: float = ( - _DBT_STATUS_REFRESH_SECONDS if refresh_seconds is None else refresh_seconds - ) - thread: threading.Thread = threading.Thread( - target=_run_active_node_status_refresher, - kwargs={ - "stream": stream, - "use_color": use_color, - "active_nodes": active_nodes, - "status_box": status_box, - "status_lock": status_lock, - "status_stop": status_stop, - "refresh_seconds": actual_refresh_seconds, - }, - daemon=True, - ) - thread.start() - return thread - - -def _run_active_node_status_refresher( - *, - stream: TextIO, - use_color: bool, - active_nodes: dict[str, tuple[str, float]], - status_box: dict[str, TransientStatusReporter | None], - status_lock: threading.Lock, - status_stop: threading.Event, - refresh_seconds: float, -) -> dict[str, TransientStatusReporter | None]: - while not status_stop.wait(refresh_seconds): - with status_lock: - status_box["status"] = _update_active_node_status( - status=status_box["status"], - active_nodes=active_nodes, - stream=stream, - use_color=use_color, - ) - return status_box - - -def _update_message_status( - *, - status: TransientStatusReporter | None, - message: str, - stream: TextIO, - use_color: bool, -) -> TransientStatusReporter | None: - if status is None: - status = TransientStatusReporter(stream=stream, use_color=use_color) - status.start(message) - return status - status.update(message) - return status - - -def _close_status( - *, - status_box: dict[str, TransientStatusReporter | None], - status_lock: threading.Lock, -) -> dict[str, TransientStatusReporter | None]: - with status_lock: - if status_box["status"] is not None: - status_box["status"].close() - status_box["status"] = None - return status_box - - -def _update_active_node_status( - *, - status: TransientStatusReporter | None, - active_nodes: dict[str, tuple[str, float]], - stream: TextIO, - use_color: bool, -) -> TransientStatusReporter | None: - if not stream.isatty() or not active_nodes: - return None - message: str = format_dbt_active_node_status( - active_nodes=list(active_nodes.values()), now=time.monotonic() - ) - if status is None: - status = TransientStatusReporter(stream=stream, use_color=use_color) - status.start(message) - return status - status.update(message) - return status - - -def _build_dbt_json_env(*, target_path: Path | None) -> dict[str, str]: - env: dict[str, str] = { - **os.environ.copy(), - "DBT_LOG_FORMAT": "json", - "PYTHONUNBUFFERED": "1", - } - if target_path is not None: - env["DBT_TARGET_PATH"] = str(target_path) - env["DBT_LOG_PATH"] = str(target_path) - return env diff --git a/src/sqlbuild/integrations/dbt/classes/__init__.py b/src/sqlbuild/integrations/dbt/classes/__init__.py deleted file mode 100644 index 0057086428..0000000000 --- a/src/sqlbuild/integrations/dbt/classes/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Public dbt integration classes.""" diff --git a/src/sqlbuild/integrations/dbt/classes/dbt_runner.py b/src/sqlbuild/integrations/dbt/classes/dbt_runner.py deleted file mode 100644 index eff420c59b..0000000000 --- a/src/sqlbuild/integrations/dbt/classes/dbt_runner.py +++ /dev/null @@ -1,125 +0,0 @@ -"""dbt CLI runner.""" - -from __future__ import annotations - -import subprocess -from collections.abc import Sequence -from pathlib import Path -from typing import cast - -from sqlbuild._native import ( - build_dbt_compile_argv, - build_dbt_debug_argv, - build_dbt_deps_argv, - build_dbt_ls_argv, - parse_dbt_ls_json_lines, - resolve_dbt_executable, -) -from sqlbuild.integrations.dbt.exceptions import DbtInteropRuntimeError -from sqlbuild.integrations.dbt.models import DbtCliOptions, DbtCommandResult, DbtLsResult -from sqlbuild.integrations.dbt.types import DbtInvoker -from sqlbuild.runtime.observability.classes.operation_lifecycle import OperationLifecycle - - -class DbtRunner: - """Small dbt CLI runner with per-invocation dbt ls memoization.""" - - def __init__( - self, *, dbt_executable: str | None = None, invoker: DbtInvoker | None = None - ) -> None: - self.dbt_executable = dbt_executable or resolve_dbt_executable() - self.invoker = invoker - self._ls_cache: dict[tuple[str, ...], DbtLsResult] = {} - - def compile(self, *, options: DbtCliOptions, full_refresh: bool = False) -> DbtCommandResult: - """Run dbt compile.""" - - argv: tuple[str, ...] = build_dbt_compile_argv( - dbt_executable=self.dbt_executable, options=options, full_refresh=full_refresh - ) - return self._invoke(argv=argv, cwd=options.project_dir) - - def deps(self, *, options: DbtCliOptions) -> DbtCommandResult: - """Run dbt deps.""" - - argv: tuple[str, ...] = build_dbt_deps_argv( - dbt_executable=self.dbt_executable, - options=options, - ) - return self._invoke(argv=argv, cwd=options.project_dir) - - def debug(self, *, options: DbtCliOptions, args: Sequence[str] = ()) -> DbtCommandResult: - """Run dbt debug.""" - - argv: tuple[str, ...] = build_dbt_debug_argv( - dbt_executable=self.dbt_executable, options=options, args=args - ) - return self._invoke(argv=argv, cwd=options.project_dir) - - def ls( - self, - *, - options: DbtCliOptions, - select: Sequence[str] = (), - exclude: Sequence[str] = (), - resource_types: Sequence[str] = (), - ) -> DbtLsResult: - """Run dbt ls and parse selected nodes.""" - - argv: tuple[str, ...] = build_dbt_ls_argv( - dbt_executable=self.dbt_executable, - options=options, - select=select, - exclude=exclude, - resource_types=resource_types, - ) - cached: DbtLsResult | None = self._ls_cache.get(argv) - if cached is not None: - return cached - command: DbtCommandResult = self._invoke(argv=argv, cwd=options.project_dir) - result: DbtLsResult = DbtLsResult( - nodes=parse_dbt_ls_json_lines(stdout=command.stdout), command=command - ) - self._ls_cache[argv] = result - return result - - def _invoke(self, *, argv: tuple[str, ...], cwd: Path | None) -> DbtCommandResult: - if self.invoker is not None: - return cast(DbtCommandResult, self.invoker(argv=argv, cwd=cwd)) - try: - with OperationLifecycle( - operation_kind="subprocess", - operation_name="dbt_command", - auto_fail_base_exceptions=False, - ) as lifecycle: - completed: subprocess.CompletedProcess[str] = subprocess.run( - argv, - cwd=cwd, - capture_output=True, - check=False, - text=True, - ) - if completed.returncode == 0: - lifecycle.completed(exit_code=0) - else: - lifecycle.failed( - error_code=f"exit_{completed.returncode}", - exit_code=completed.returncode, - signal_number=(-completed.returncode if completed.returncode < 0 else None), - ) - except OSError as error: - raise DbtInteropRuntimeError( - "failed to execute dbt", - help=str(error), - ) from error - return DbtCommandResult( - argv=argv, - returncode=completed.returncode, - stdout=completed.stdout, - stderr=completed.stderr, - ) - - def invoke(self, *, argv: tuple[str, ...], cwd: Path | None) -> DbtCommandResult: - """Run an explicit dbt argv.""" - - return self._invoke(argv=argv, cwd=cwd) diff --git a/src/sqlbuild/integrations/dbt/main/runtime/__init__.py b/src/sqlbuild/integrations/dbt/main/runtime/__init__.py deleted file mode 100644 index 04ef776caf..0000000000 --- a/src/sqlbuild/integrations/dbt/main/runtime/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""dbt runtime public operations.""" diff --git a/src/sqlbuild/integrations/dbt/main/runtime/_execute_json_event_stream.py b/src/sqlbuild/integrations/dbt/main/runtime/_execute_json_event_stream.py deleted file mode 100644 index 3498826808..0000000000 --- a/src/sqlbuild/integrations/dbt/main/runtime/_execute_json_event_stream.py +++ /dev/null @@ -1,37 +0,0 @@ -"""Execute a dbt JSON event stream.""" - -from collections.abc import Callable -from pathlib import Path -from typing import TextIO - -from sqlbuild.integrations.dbt._helpers.runtime.event_stream import ( - execute_dbt_json_event_stream as _execute, -) -from sqlbuild.integrations.dbt.models import DbtNodeExecutionResult - - -def execute_dbt_json_event_stream( - *, - argv: tuple[str, ...], - cwd: Path | None, - stream: TextIO, - use_color: bool, - target_path: Path | None, - display_total: int | None = None, - on_node_result: Callable[[DbtNodeExecutionResult], None] | None = None, - detail_by_unique_id: dict[str, str] | None = None, - enable_status: bool = True, -) -> tuple[int, tuple[DbtNodeExecutionResult, ...]]: - """Run dbt and render SQLBuild-styled JSON events.""" - - return _execute( - argv=argv, - cwd=cwd, - stream=stream, - use_color=use_color, - target_path=target_path, - display_total=display_total, - on_node_result=on_node_result, - detail_by_unique_id=detail_by_unique_id, - enable_status=enable_status, - ) diff --git a/tests/integration/src/sqlbuild/integrations/dbt/test_graph.py b/tests/integration/src/sqlbuild/integrations/dbt/test_graph.py index 4cd962a4ce..8a4d126f02 100644 --- a/tests/integration/src/sqlbuild/integrations/dbt/test_graph.py +++ b/tests/integration/src/sqlbuild/integrations/dbt/test_graph.py @@ -8,6 +8,7 @@ from sqlbuild._native import ( DbtCombinedGraph, DbtManifestIndex, + DbtRunner, assemble_project, build_compile_inputs, build_dbt_combined_graph, @@ -16,7 +17,6 @@ ) from sqlbuild.compiler.compile.types import CompiledProjectView, CompileInputsView from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs -from sqlbuild.integrations.dbt.classes.dbt_runner import DbtRunner from sqlbuild.integrations.dbt.models import DbtCliOptions, DbtCommandResult from tests.integration.src.sqlbuild.integrations.dbt._test_types import ( RealDbtCombinedGraphTestCase, diff --git a/tests/integration/src/sqlbuild/integrations/dbt/test_manifest.py b/tests/integration/src/sqlbuild/integrations/dbt/test_manifest.py index 27d3dc8921..6b64d7125a 100644 --- a/tests/integration/src/sqlbuild/integrations/dbt/test_manifest.py +++ b/tests/integration/src/sqlbuild/integrations/dbt/test_manifest.py @@ -4,10 +4,9 @@ import pytest -from sqlbuild._native import build_compile_inputs, discover_project_inputs +from sqlbuild._native import DbtRunner, build_compile_inputs, discover_project_inputs from sqlbuild.compiler.compile.types import CompileInputsView from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs -from sqlbuild.integrations.dbt.classes.dbt_runner import DbtRunner from sqlbuild.integrations.dbt.models import DbtCliOptions, DbtCommandResult from tests.integration.src.sqlbuild.integrations.dbt._test_types import ( RealDbtManifestCompileTestCase, diff --git a/tests/integration/src/sqlbuild/integrations/dbt/test_runner.py b/tests/integration/src/sqlbuild/integrations/dbt/test_runner.py index b854238cd4..f98446c2d8 100644 --- a/tests/integration/src/sqlbuild/integrations/dbt/test_runner.py +++ b/tests/integration/src/sqlbuild/integrations/dbt/test_runner.py @@ -4,7 +4,7 @@ import pytest -from sqlbuild.integrations.dbt.classes.dbt_runner import DbtRunner +from sqlbuild._native import DbtRunner from sqlbuild.integrations.dbt.models import DbtCliOptions, DbtCommandResult, DbtLsResult from tests.integration.src.sqlbuild.integrations.dbt._test_types import RealDbtRunnerTestCase diff --git a/tests/unit/src/sqlbuild/cli/commands/main/dbt/test_dbt.py b/tests/unit/src/sqlbuild/cli/commands/main/dbt/test_dbt.py index 737749767a..a7a6672920 100644 --- a/tests/unit/src/sqlbuild/cli/commands/main/dbt/test_dbt.py +++ b/tests/unit/src/sqlbuild/cli/commands/main/dbt/test_dbt.py @@ -288,7 +288,7 @@ def plan_dbt_interop_from_project( return build_empty_dbt_plan() monkeypatch.setattr( - "sqlbuild.integrations.dbt._helpers.profile.init.build_dbt_init_project", + "sqlbuild._native.build_dbt_init_project", run_dbt_profile_init, ) monkeypatch.setattr( @@ -377,7 +377,7 @@ def plan_dbt_interop_from_project( return build_empty_dbt_plan() monkeypatch.setattr( - "sqlbuild.integrations.dbt._helpers.profile.init.build_dbt_init_project", + "sqlbuild._native.build_dbt_init_project", run_dbt_profile_init, ) monkeypatch.setattr( diff --git a/tests/unit/src/sqlbuild/cli/commands/main/dbt/test_init.py b/tests/unit/src/sqlbuild/cli/commands/main/dbt/test_init.py index 9325a2e1c3..df806bc7b3 100644 --- a/tests/unit/src/sqlbuild/cli/commands/main/dbt/test_init.py +++ b/tests/unit/src/sqlbuild/cli/commands/main/dbt/test_init.py @@ -6,7 +6,7 @@ from _pytest.capture import CaptureResult import sqlbuild._native as dbt_init_module -import sqlbuild.integrations.dbt._helpers.profile.init as dbt_profile_init +import sqlbuild._native as dbt_profile_init from sqlbuild._native import DbtInitCommandRequest from sqlbuild.integrations.dbt.exceptions import DbtProfileError from sqlbuild.integrations.dbt.models import DbtInitRequest, DbtInitResult diff --git a/tests/unit/src/sqlbuild/integrations/dbt/_test_types.py b/tests/unit/src/sqlbuild/integrations/dbt/_test_types.py index 2c4faefd40..0a49893e90 100644 --- a/tests/unit/src/sqlbuild/integrations/dbt/_test_types.py +++ b/tests/unit/src/sqlbuild/integrations/dbt/_test_types.py @@ -334,3 +334,19 @@ class DbtPlanOrchestrationErrorTestCase: expected_error_fragment: str expected_code: str expected_help_fragment: str + + +@dataclass(frozen=True) +class DbtRunnerProcessTestCase: + description: str + script: str + expected_returncode: int + expected_stdout: str + expected_stderr: str + + +@dataclass(frozen=True) +class DbtRunnerLaunchFailureTestCase: + description: str + executable_name: str + expected_message: str diff --git a/tests/unit/src/sqlbuild/integrations/dbt/helpers.py b/tests/unit/src/sqlbuild/integrations/dbt/helpers.py index 46e4b5ded1..d11d0408dc 100644 --- a/tests/unit/src/sqlbuild/integrations/dbt/helpers.py +++ b/tests/unit/src/sqlbuild/integrations/dbt/helpers.py @@ -2,6 +2,7 @@ import json import re +import sys from collections.abc import Callable from pathlib import Path from types import MappingProxyType @@ -13,6 +14,7 @@ DbtCombinedGraph, DbtEventStream, DbtEventStreamUpdate, + DbtRunner, assemble_project, build_dbt_combined_graph, build_dbt_ls_argv, @@ -28,13 +30,8 @@ CompileProjectInputs, ) from sqlbuild.compiler.compile.types import CompiledProjectView, CompiledResourceType -from sqlbuild.compiler.discovery.models import ( - DiscoveredProjectInputs, - DiscoveredSqlModelFile, -) -from sqlbuild.compiler.planner.models import ( - ModelPlanEntry, -) +from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs, DiscoveredSqlModelFile +from sqlbuild.compiler.planner.models import ModelPlanEntry from sqlbuild.compiler.planner.types import ( BackfillAction, MaterializationType, @@ -42,7 +39,6 @@ PlanReason, ) from sqlbuild.compiler.sql_analysis.models import SqlLexicalSyntax -from sqlbuild.integrations.dbt.classes.dbt_runner import DbtRunner from sqlbuild.integrations.dbt.models import DbtCliOptions, DbtCommandResult from sqlbuild.spec.contracts.models import LocalConfig, ProjectConfig @@ -495,3 +491,24 @@ def consume_dbt_events(*events: dict[str, object]) -> DbtEventStreamUpdate | Non stream.consume(json.dumps(event)) for event in events ) return updates[-1] + + +def fake_dbt_argv( + directory: Path, + *, + lines: tuple[str, ...], + sleep_seconds: float = 0.0, + after_lines: tuple[str, ...] = (), +) -> tuple[str, ...]: + """argv of a stand-in dbt process printing `lines`, pausing, then printing `after_lines`.""" + + script: Path = directory / "fake_dbt.py" + _ = script.write_text( + "import sys, time\n" + f"sys.stdout.write({''.join(lines)!r})\n" + "sys.stdout.flush()\n" + f"time.sleep({sleep_seconds!r})\n" + f"sys.stdout.write({''.join(after_lines)!r})\n", + encoding="utf-8", + ) + return (sys.executable, str(script)) diff --git a/tests/unit/src/sqlbuild/integrations/dbt/profile/test_init.py b/tests/unit/src/sqlbuild/integrations/dbt/profile/test_init.py index 861c8a2390..985c8bf50c 100644 --- a/tests/unit/src/sqlbuild/integrations/dbt/profile/test_init.py +++ b/tests/unit/src/sqlbuild/integrations/dbt/profile/test_init.py @@ -4,7 +4,7 @@ import pytest -from sqlbuild.integrations.dbt._helpers.profile.init import build_dbt_init_project +from sqlbuild._native import build_dbt_init_project from sqlbuild.integrations.dbt.models import DbtInitRequest, DbtInitResult from tests.unit.src.sqlbuild.integrations.dbt.profile._test_types import ( DbtProfileInitDiscoveryTestCase, diff --git a/tests/unit/src/sqlbuild/integrations/dbt/test_debug.py b/tests/unit/src/sqlbuild/integrations/dbt/test_debug.py index b1b3b2e14d..3f76d490da 100644 --- a/tests/unit/src/sqlbuild/integrations/dbt/test_debug.py +++ b/tests/unit/src/sqlbuild/integrations/dbt/test_debug.py @@ -5,8 +5,7 @@ import pytest -from sqlbuild._native import debug_dbt_from_project -from sqlbuild.integrations.dbt.classes.dbt_runner import DbtRunner +from sqlbuild._native import DbtRunner, debug_dbt_from_project from sqlbuild.integrations.dbt.models import DbtCommandResult from tests.unit.src.sqlbuild.integrations.dbt._test_types import DbtDebugPipelineTestCase from tests.unit.src.sqlbuild.integrations.dbt.helpers import RecordingDbtInvoker diff --git a/tests/unit/src/sqlbuild/integrations/dbt/test_event_stream.py b/tests/unit/src/sqlbuild/integrations/dbt/test_event_stream.py index eec7cdfffe..4fb1ecf21c 100644 --- a/tests/unit/src/sqlbuild/integrations/dbt/test_event_stream.py +++ b/tests/unit/src/sqlbuild/integrations/dbt/test_event_stream.py @@ -2,17 +2,17 @@ import io import json -import threading import time -from collections.abc import Callable, Iterator -from typing import cast +from pathlib import Path import pytest -from sqlbuild._native import DbtEventStream, DbtEventStreamUpdate -from sqlbuild.integrations.dbt._helpers.runtime.event_stream import ( +from sqlbuild._native import ( + DbtEventStream, + DbtEventStreamUpdate, execute_dbt_json_event_stream, ) +from sqlbuild.integrations.dbt.exceptions import DbtInteropRuntimeError from sqlbuild.integrations.dbt.models import DbtNodeExecutionResult from sqlbuild.observability import ( EventDispatcher, @@ -26,8 +26,20 @@ DbtRuntimeCleanupTestCase, DbtSilentStatusRefreshTestCase, ) -from tests.unit.src.sqlbuild.integrations.dbt.helpers import consume_dbt_events +from tests.unit.src.sqlbuild.integrations.dbt.helpers import consume_dbt_events, fake_dbt_argv +_STOPPED_WITHIN_SECONDS: float = 30.0 +_DBT_SLOW_START_LINE: str = ( + '{"data":{"index":1,"total":1,"node_info":{"node_name":"slow_model",' + '"resource_type":"model","unique_id":"model.analytics.slow_model"}},' + '"info":{"level":"info","name":"LogStartLine","msg":"START"}}\n' +) +_DBT_SLOW_RESULT_LINE: str = ( + '{"data":{"execution_time":1.2,"index":1,"total":1,"status":"success",' + '"node_info":{"node_name":"slow_model","resource_type":"model",' + '"node_status":"success","unique_id":"model.analytics.slow_model"}},' + '"info":{"level":"info","name":"LogModelResult","msg":"OK"}}\n' +) _DBT_RESULT_LINE: str = ( '{"data":{"execution_time":1.0,"index":1,"total":1,"status":"success",' '"node_info":{"node_name":"orders","resource_type":"model",' @@ -409,35 +421,17 @@ def test_given_non_result_event_when_parsing_then_returns_none( ) def test_given_dbt_json_stream_when_running_then_invokes_node_result_callback( test_case: DbtEventStreamTestCase, - monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, ) -> None: lifecycle_events: list[LifecycleEvent] = [] - - class StubProcess: - stdout: io.StringIO = io.StringIO("".join(test_case.stdout_lines)) - - def wait(self) -> int: - assert tuple(event.event_type for event in lifecycle_events) == ("operation_started",) - return 0 - captured_results: list[DbtNodeExecutionResult] = [] - - def launch(*args: object, **kwargs: object) -> StubProcess: - del args, kwargs - assert lifecycle_events[-1].event_type == "operation_started" - return StubProcess() - - monkeypatch.setattr( - "sqlbuild.integrations.dbt._helpers.runtime.event_stream.subprocess.Popen", launch - ) - stream: io.StringIO = io.StringIO() dispatcher: EventDispatcher = EventDispatcher() dispatcher.subscribe_lifecycle(subscriber=lifecycle_events.append, accepts_opaque=False) - returncode: int + with invocation_scope("inv-dbt-stream"), dispatcher_scope(dispatcher): returncode, results = execute_dbt_json_event_stream( - argv=("dbt", "run"), + argv=fake_dbt_argv(tmp_path, lines=test_case.stdout_lines), cwd=None, stream=stream, use_color=False, @@ -467,101 +461,23 @@ def launch(*args: object, **kwargs: object) -> StubProcess: "test_case", ( DbtRuntimeCleanupTestCase( - description="callback failure cleans launched dbt runtime before terminal", + description="callback failure stops dbt before the failed terminal", expected_error="original callback failure", - expected_actions=( - "operation_started", - "popen", - "status_start", - "thread_start", - "status_close", - "thread_join", - "poll", - "terminate", - "wait", - "operation_failed", - ), + expected_actions=("operation_started", "operation_failed"), ), ), ids=lambda case: case.description, ) def test_given_callback_failure_when_streaming_then_runtime_cleans_before_failed_terminal( test_case: DbtRuntimeCleanupTestCase, - monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, ) -> None: actions: list[str] = [] - lifecycle_events: list[LifecycleEvent] = [] - - class TtyStream(io.StringIO): - def isatty(self) -> bool: - return True - - class StatusReporter: - def __init__(self, *args: object, **kwargs: object) -> None: - del args, kwargs - - def start(self, message: str) -> None: - del message - actions.append("status_start") - - def close(self) -> None: - assert lifecycle_events[-1].event_type == "operation_started" - actions.append("status_close") - - class StatusThread: - def __init__(self, status_stop: threading.Event) -> None: - self._status_stop: threading.Event = status_stop - - def join(self, timeout: float) -> None: - del timeout - assert lifecycle_events[-1].event_type == "operation_started" - assert self._status_stop.is_set() - actions.append("thread_join") - - class Process: - stdout: io.StringIO = io.StringIO(_DBT_RESULT_LINE) - - def poll(self) -> None: - actions.append("poll") - return None - - def terminate(self) -> None: - assert lifecycle_events[-1].event_type == "operation_started" - actions.append("terminate") - - def wait(self, timeout: float | None = None) -> int: - assert timeout == 1 - assert lifecycle_events[-1].event_type == "operation_started" - actions.append("wait") - return -15 - - def launch(*args: object, **kwargs: object) -> Process: - del args, kwargs - actions.append("popen") - return Process() - - def start_thread(*args: object, **kwargs: object) -> StatusThread: - del args - actions.append("thread_start") - return StatusThread(cast(threading.Event, kwargs["status_stop"])) - - def record_event(event: LifecycleEvent) -> None: - lifecycle_events.append(event) - actions.append(event.event_type) - - monkeypatch.setattr( - "sqlbuild.integrations.dbt._helpers.runtime.event_stream.subprocess.Popen", launch - ) - monkeypatch.setattr( - "sqlbuild.integrations.dbt._helpers.runtime.event_stream.TransientStatusReporter", - StatusReporter, - ) - monkeypatch.setattr( - "sqlbuild.integrations.dbt._helpers.runtime.event_stream._start_active_node_status_refresher", - start_thread, - ) dispatcher: EventDispatcher = EventDispatcher() - dispatcher.subscribe_lifecycle(subscriber=record_event, accepts_opaque=False) + dispatcher.subscribe_lifecycle( + subscriber=lambda event: actions.append(event.event_type), accepts_opaque=False + ) + started: float = time.monotonic() with ( invocation_scope("inv-dbt-callback-failure"), @@ -569,9 +485,9 @@ def record_event(event: LifecycleEvent) -> None: pytest.raises(RuntimeError, match=test_case.expected_error), ): _ = execute_dbt_json_event_stream( - argv=("dbt", "run"), + argv=fake_dbt_argv(tmp_path, lines=(_DBT_RESULT_LINE,), sleep_seconds=60), cwd=None, - stream=TtyStream(), + stream=io.StringIO(), use_color=False, target_path=None, on_node_result=lambda result: (_ for _ in ()).throw( @@ -580,295 +496,47 @@ def record_event(event: LifecycleEvent) -> None: ) assert tuple(actions) == test_case.expected_actions + assert time.monotonic() - started < _STOPPED_WITHIN_SECONDS @pytest.mark.parametrize( "test_case", ( DbtRuntimeCleanupTestCase( - description="process wait failure reaps child before terminal", - expected_error="original wait failure", - expected_actions=( - "operation_started", - "main_wait", - "poll", - "terminate", - "cleanup_wait", - "operation_failed", - ), - ), - ), - ids=lambda case: case.description, -) -def test_given_process_wait_failure_when_streaming_then_child_cleanup_precedes_terminal( - test_case: DbtRuntimeCleanupTestCase, - monkeypatch: pytest.MonkeyPatch, -) -> None: - actions: list[str] = [] - lifecycle_events: list[LifecycleEvent] = [] - wait_labels: Iterator[str] = iter(("main_wait", "cleanup_wait")) - wait_outcomes: Iterator[Callable[[], int]] = iter( - ( - lambda: (_ for _ in ()).throw(RuntimeError(test_case.expected_error)), - lambda: -15, - ) - ) - - class Process: - stdout: io.StringIO = io.StringIO() - - def poll(self) -> None: - actions.append("poll") - return None - - def terminate(self) -> None: - assert lifecycle_events[-1].event_type == "operation_started" - actions.append("terminate") - - def wait(self, timeout: float | None = None) -> int: - del timeout - assert lifecycle_events[-1].event_type == "operation_started" - actions.append(next(wait_labels)) - return next(wait_outcomes)() - - def record_event(event: LifecycleEvent) -> None: - lifecycle_events.append(event) - actions.append(event.event_type) - - monkeypatch.setattr( - "sqlbuild.integrations.dbt._helpers.runtime.event_stream.subprocess.Popen", - lambda *args, **kwargs: Process(), - ) - dispatcher: EventDispatcher = EventDispatcher() - dispatcher.subscribe_lifecycle(subscriber=record_event, accepts_opaque=False) - - with ( - invocation_scope("inv-dbt-wait-failure"), - dispatcher_scope(dispatcher), - pytest.raises(RuntimeError, match=test_case.expected_error), - ): - _ = execute_dbt_json_event_stream( - argv=("dbt", "run"), - cwd=None, - stream=io.StringIO(), - use_color=False, - target_path=None, - enable_status=False, - ) - - assert tuple(actions) == test_case.expected_actions - - -@pytest.mark.parametrize( - "test_case", - ( - DbtRuntimeCleanupTestCase( - description="popen failure starts no status thread", + description="launch failure starts no status and fails the operation", expected_error="failed to execute dbt", - expected_actions=("operation_started", "popen", "operation_failed"), + expected_actions=("operation_started", "operation_failed"), ), ), ids=lambda case: case.description, ) def test_given_popen_failure_when_streaming_then_no_status_runtime_is_started( test_case: DbtRuntimeCleanupTestCase, - monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, ) -> None: actions: list[str] = [] dispatcher: EventDispatcher = EventDispatcher() dispatcher.subscribe_lifecycle( subscriber=lambda event: actions.append(event.event_type), accepts_opaque=False ) - - def launch(*args: object, **kwargs: object) -> None: - del args, kwargs - actions.append("popen") - raise OSError("private launch failure") - - monkeypatch.setattr( - "sqlbuild.integrations.dbt._helpers.runtime.event_stream.subprocess.Popen", launch - ) - monkeypatch.setattr( - "sqlbuild.integrations.dbt._helpers.runtime.event_stream._start_active_node_status_refresher", - lambda *args, **kwargs: actions.append("unexpected_thread"), - ) + stream: io.StringIO = io.StringIO() with ( invocation_scope("inv-dbt-popen-failure"), dispatcher_scope(dispatcher), - pytest.raises(Exception, match=test_case.expected_error), + pytest.raises(DbtInteropRuntimeError, match=test_case.expected_error) as raised, ): _ = execute_dbt_json_event_stream( - argv=("dbt", "run"), + argv=(str(tmp_path / "missing-dbt"), "run"), cwd=None, - stream=io.StringIO(), - use_color=False, - target_path=None, - ) - - assert tuple(actions) == test_case.expected_actions - - -@pytest.mark.parametrize( - "test_case", - ( - DbtRuntimeCleanupTestCase( - description="interrupted dbt wait leaves operation unmatched", - expected_error="", - expected_actions=("operation_started",), - ), - ), - ids=lambda case: case.description, -) -def test_given_unobserved_dbt_interruption_when_cleanup_reaps_then_no_terminal_is_fabricated( - test_case: DbtRuntimeCleanupTestCase, - monkeypatch: pytest.MonkeyPatch, -) -> None: - lifecycle_events: list[LifecycleEvent] = [] - wait_outcomes: Iterator[Callable[[], int]] = iter( - ( - lambda: (_ for _ in ()).throw(KeyboardInterrupt), - lambda: -2, - ) - ) - - class Process: - stdout: io.StringIO = io.StringIO() - - def poll(self) -> int: - return -2 - - def wait(self, timeout: float | None = None) -> int: - del timeout - return next(wait_outcomes)() - - monkeypatch.setattr( - "sqlbuild.integrations.dbt._helpers.runtime.event_stream.subprocess.Popen", - lambda *args, **kwargs: Process(), - ) - dispatcher: EventDispatcher = EventDispatcher() - dispatcher.subscribe_lifecycle(subscriber=lifecycle_events.append, accepts_opaque=False) - - with ( - invocation_scope("inv-dbt-interrupted"), - dispatcher_scope(dispatcher), - pytest.raises(KeyboardInterrupt), - ): - _ = execute_dbt_json_event_stream( - argv=("dbt", "run"), - cwd=None, - stream=io.StringIO(), - use_color=False, - target_path=None, - enable_status=False, - ) - - assert tuple(event.event_type for event in lifecycle_events) == test_case.expected_actions - - -@pytest.mark.parametrize( - "test_case", - ( - DbtRuntimeCleanupTestCase( - description="cleanup failures preserve parse exception", - expected_error="original parse failure", - expected_actions=( - "operation_started", - "thread_join", - "status_close", - "poll", - "terminate", - "wait", - "kill", - "wait", - "operation_failed", - ), - ), - ), - ids=lambda case: case.description, -) -def test_given_cleanup_failures_when_parse_raises_then_original_exception_is_preserved( - test_case: DbtRuntimeCleanupTestCase, - monkeypatch: pytest.MonkeyPatch, -) -> None: - actions: list[str] = [] - - class TtyStream(io.StringIO): - def isatty(self) -> bool: - return True - - class FailingStatus: - def __init__(self, *args: object, **kwargs: object) -> None: - del args, kwargs - - def start(self, message: str) -> None: - del message - - def close(self) -> None: - actions.append("status_close") - raise RuntimeError("cleanup status failure") - - class FailingThread: - def join(self, timeout: float) -> None: - del timeout - actions.append("thread_join") - raise RuntimeError("cleanup thread failure") - - class FailingProcess: - stdout: io.StringIO = io.StringIO(_DBT_RESULT_LINE) - - def poll(self) -> None: - actions.append("poll") - raise RuntimeError("cleanup poll failure") - - def terminate(self) -> None: - actions.append("terminate") - raise RuntimeError("cleanup terminate failure") - - def wait(self, timeout: float | None = None) -> int: - del timeout - actions.append("wait") - raise RuntimeError("cleanup wait failure") - - def kill(self) -> None: - actions.append("kill") - raise RuntimeError("cleanup kill failure") - - dispatcher: EventDispatcher = EventDispatcher() - dispatcher.subscribe_lifecycle( - subscriber=lambda event: actions.append(event.event_type), accepts_opaque=False - ) - monkeypatch.setattr( - "sqlbuild.integrations.dbt._helpers.runtime.event_stream.subprocess.Popen", - lambda *args, **kwargs: FailingProcess(), - ) - monkeypatch.setattr( - "sqlbuild.integrations.dbt._helpers.runtime.event_stream.TransientStatusReporter", - FailingStatus, - ) - monkeypatch.setattr( - "sqlbuild.integrations.dbt._helpers.runtime.event_stream._start_active_node_status_refresher", - lambda *args, **kwargs: FailingThread(), - ) - monkeypatch.setattr( - "sqlbuild.integrations.dbt._helpers.runtime.event_stream._handle_dbt_output_line", - lambda **kwargs: (_ for _ in ()).throw(RuntimeError(test_case.expected_error)), - ) - - with ( - invocation_scope("inv-dbt-cleanup-failure"), - dispatcher_scope(dispatcher), - pytest.raises(RuntimeError, match=test_case.expected_error), - ): - _ = execute_dbt_json_event_stream( - argv=("dbt", "run"), - cwd=None, - stream=TtyStream(), + stream=stream, use_color=False, target_path=None, ) assert tuple(actions) == test_case.expected_actions + assert isinstance(raised.value.__cause__, FileNotFoundError) + assert stream.getvalue() == "" @pytest.mark.parametrize( @@ -876,8 +544,8 @@ def kill(self) -> None: [ DbtSilentStatusRefreshTestCase( description="active node elapsed refreshes without dbt events", - silent_seconds=1.15, - refresh_seconds=0.05, + silent_seconds=2.4, + refresh_seconds=1.0, expected_initial_status="running 1 dbt node: slow_model <1s", expected_refreshed_status="running 1 dbt node: slow_model 1s", ) @@ -886,53 +554,23 @@ def kill(self) -> None: ) def test_given_active_dbt_node_when_dbt_stream_is_silent_then_status_elapsed_updates( test_case: DbtSilentStatusRefreshTestCase, + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: class TtyStream(io.StringIO): def isatty(self) -> bool: return True - class DelayedStdout: - def __enter__(self) -> DelayedStdout: - return self - - def __exit__(self, *args: object) -> None: - return None - - def __iter__(self) -> object: - yield ( - '{"data":{"index":1,"total":1,"node_info":{"node_name":"slow_model",' - '"resource_type":"model","unique_id":"model.analytics.slow_model"}},' - '"info":{"level":"info","name":"LogStartLine","msg":"START"}}\n' - ) - time.sleep(test_case.silent_seconds) - yield ( - '{"data":{"execution_time":1.2,"index":1,"total":1,"status":"success",' - '"node_info":{"node_name":"slow_model","resource_type":"model",' - '"node_status":"success","unique_id":"model.analytics.slow_model"}},' - '"info":{"level":"info","name":"LogModelResult","msg":"OK"}}\n' - ) - - class StubProcess: - stdout: DelayedStdout = DelayedStdout() - - def wait(self) -> int: - return 0 - monkeypatch.setenv("SQLBUILD_NO_PROGRESS", "1") - monkeypatch.setattr( - "sqlbuild.integrations.dbt._helpers.runtime.event_stream._DBT_STATUS_REFRESH_SECONDS", - test_case.refresh_seconds, - ) - monkeypatch.setattr( - "sqlbuild.integrations.dbt._helpers.runtime.event_stream.subprocess.Popen", - lambda *args, **kwargs: StubProcess(), - ) - stream: TtyStream = TtyStream() - returncode: int + returncode, _ = execute_dbt_json_event_stream( - argv=("dbt", "run"), + argv=fake_dbt_argv( + tmp_path, + lines=(_DBT_SLOW_START_LINE,), + sleep_seconds=test_case.silent_seconds, + after_lines=(_DBT_SLOW_RESULT_LINE,), + ), cwd=None, stream=stream, use_color=False, @@ -941,6 +579,7 @@ def wait(self) -> int: status_lines: list[str] = stream.getvalue().splitlines() assert returncode == 0 + assert test_case.refresh_seconds < test_case.silent_seconds assert test_case.expected_initial_status in status_lines assert test_case.expected_refreshed_status in status_lines diff --git a/tests/unit/src/sqlbuild/integrations/dbt/test_execute.py b/tests/unit/src/sqlbuild/integrations/dbt/test_execute.py index 261c56e695..4a6051fd9f 100644 --- a/tests/unit/src/sqlbuild/integrations/dbt/test_execute.py +++ b/tests/unit/src/sqlbuild/integrations/dbt/test_execute.py @@ -8,6 +8,7 @@ import pytest +import sqlbuild._native as event_module from sqlbuild import _native as execute_module from sqlbuild import _native as phases_module from sqlbuild import _native as plan_module @@ -24,6 +25,7 @@ DbtInteropSelectionResult, DbtLsNode, DbtManifestIndex, + DbtRunner, build_dbt_interop_plan, build_failed_sqlbuild_model_names, build_merged_dbt_execution_argv, @@ -32,8 +34,6 @@ from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.compiler.compile.models import CompiledProject from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs -from sqlbuild.integrations.dbt.classes.dbt_runner import DbtRunner -from sqlbuild.integrations.dbt.main.runtime import _execute_json_event_stream as event_module from sqlbuild.integrations.dbt.models import DbtCliOptions, DbtNodeExecutionResult from tests.unit.src.sqlbuild.integrations.dbt._test_types import ( DbtCompileFullRefreshPipelineTestCase, diff --git a/tests/unit/src/sqlbuild/integrations/dbt/test_plan_orchestration.py b/tests/unit/src/sqlbuild/integrations/dbt/test_plan_orchestration.py index 51e6793cc8..60aa3dd288 100644 --- a/tests/unit/src/sqlbuild/integrations/dbt/test_plan_orchestration.py +++ b/tests/unit/src/sqlbuild/integrations/dbt/test_plan_orchestration.py @@ -9,13 +9,13 @@ DbtInteropCommandArgs, DbtInteropPlan, DbtManifestIndex, + DbtRunner, build_dbt_combined_graph, build_dbt_ls_argv, build_dbt_manifest_index, plan_dbt_interop_command, ) from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.integrations.dbt.classes.dbt_runner import DbtRunner from sqlbuild.integrations.dbt.exceptions import DbtInteropRuntimeError from sqlbuild.integrations.dbt.models import DbtCliOptions, DbtCommandResult from tests.unit.src.sqlbuild.integrations.dbt._test_types import ( diff --git a/tests/unit/src/sqlbuild/integrations/dbt/test_runner.py b/tests/unit/src/sqlbuild/integrations/dbt/test_runner.py index e28c248bf8..0aad6d5cb8 100644 --- a/tests/unit/src/sqlbuild/integrations/dbt/test_runner.py +++ b/tests/unit/src/sqlbuild/integrations/dbt/test_runner.py @@ -1,11 +1,13 @@ from __future__ import annotations +import sys from pathlib import Path import pytest from sqlbuild._native import ( DbtLsNode, + DbtRunner, build_dbt_compile_argv, build_dbt_debug_argv, build_dbt_deps_argv, @@ -13,13 +15,19 @@ dbt_fqn_selector_term, parse_dbt_ls_json_lines, ) -from sqlbuild.integrations.dbt.classes.dbt_runner import DbtRunner -from sqlbuild.integrations.dbt.models import DbtCliOptions, DbtCommandResult, DbtLsResult +from sqlbuild.integrations.dbt.exceptions import DbtInteropRuntimeError +from sqlbuild.integrations.dbt.models import ( + DbtCliOptions, + DbtCommandResult, + DbtLsResult, +) from tests.unit.src.sqlbuild.integrations.dbt._test_types import ( DbtArgvTestCase, DbtLsParseTestCase, DbtRunnerCommandTestCase, + DbtRunnerLaunchFailureTestCase, DbtRunnerMemoTestCase, + DbtRunnerProcessTestCase, ) from tests.unit.src.sqlbuild.integrations.dbt.helpers import ( RecordingDbtInvoker, @@ -296,3 +304,56 @@ def test_given_dbt_runner_when_running_compile_then_uses_project_dir_as_cwd( assert result == test_case.command_result assert invoker.calls == [(test_case.expected_argv, options.project_dir)] + + +@pytest.mark.parametrize( + "test_case", + [ + DbtRunnerProcessTestCase( + description="captures text output, stderr and the exit code of a real process", + script="import sys\nprint('compiled')\nsys.stderr.write('warn\\r\\n')\nsys.exit(3)\n", + expected_returncode=3, + expected_stdout="compiled\n", + expected_stderr="warn\n", + ) + ], + ids=lambda case: case.description, +) +def test_given_dbt_process_when_invoked_then_captures_output_and_exit_code( + test_case: DbtRunnerProcessTestCase, + tmp_path: Path, +) -> None: + script: Path = tmp_path / "fake_dbt.py" + _ = script.write_text(test_case.script, encoding="utf-8") + runner: DbtRunner = DbtRunner(dbt_executable="dbt") + + result: DbtCommandResult = runner.invoke(argv=(sys.executable, str(script)), cwd=tmp_path) + + assert result.returncode == test_case.expected_returncode + assert result.stdout == test_case.expected_stdout + assert result.stderr == test_case.expected_stderr + + +@pytest.mark.parametrize( + "test_case", + [ + DbtRunnerLaunchFailureTestCase( + description="a missing executable raises the dbt launch error", + executable_name="missing-dbt", + expected_message="failed to execute dbt", + ) + ], + ids=lambda case: case.description, +) +def test_given_missing_dbt_executable_when_invoked_then_raises_launch_error( + test_case: DbtRunnerLaunchFailureTestCase, + tmp_path: Path, +) -> None: + runner: DbtRunner = DbtRunner(dbt_executable="dbt") + executable: Path = tmp_path / test_case.executable_name + + with pytest.raises(DbtInteropRuntimeError, match=test_case.expected_message) as raised: + _ = runner.invoke(argv=(str(executable),), cwd=None) + + assert isinstance(raised.value.__cause__, FileNotFoundError) + assert raised.value.help == f"[Errno 2] No such file or directory: '{executable}'" From 9ce78285aaaf4cb468f71852b84ce7eaed13e56b Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:27:54 +0100 Subject: [PATCH 03/12] refactor: orchestrate custom rules and their host processes natively --- .../bindings/_helpers/cache/native_store.rs | 16 +- .../_helpers/compile_driver/analysis.rs | 2 +- .../_helpers/compile_driver/pipeline.rs | 2 +- .../_helpers/dbt_interop/event_stream.rs | 25 +- .../_helpers/rules/custom_evaluation.rs | 495 ++++++++++++++ .../bindings/_helpers/rules/custom_hosts.rs | 399 +++++++++++ .../_helpers/rules/custom_rule_cache.rs | 72 +- .../bindings/_helpers/rules/early_sql_lint.rs | 2 +- .../src/bindings/_helpers/rules/mod.rs | 3 + .../bindings/_helpers/rules/registration.rs | 3 +- .../bindings/_helpers/rules/rule_dispatch.rs | 621 ++++++++++++++++++ .../src/bindings/_helpers/rules/rule_run.rs | 3 +- .../src/custom_host_pool/_helpers/cgroups.rs | 83 +++ .../src/custom_host_pool/_helpers/decoding.rs | 119 ++++ .../src/custom_host_pool/_helpers/merging.rs | 116 ++++ .../src/custom_host_pool/_helpers/mod.rs | 4 + .../_helpers/payload_values.rs | 36 + .../src/custom_host_pool/constants.rs | 27 + .../custom_host_pool/main/available_cores.rs | 17 + .../main/decode_host_response.rs | 11 + .../src/custom_host_pool/main/host_count.rs | 12 + .../main/merge_host_slices.rs | 32 + .../src/custom_host_pool/main/mod.rs | 7 + .../main/partition_host_plan.rs | 47 ++ .../main/planned_invocations.rs | 40 ++ .../custom_host_pool/main/without_codes.rs | 25 + .../src/custom_host_pool/mod.rs | 10 + .../src/custom_host_pool/models.rs | 42 ++ .../src/custom_host_pool/tests.rs | 6 + .../src/custom_host_pool/tests/helpers.rs | 56 ++ .../tests/test_custom_host_pool.rs | 294 +++++++++ .../src/custom_host_pool/tests/test_types.rs | 39 ++ .../src/custom_host_pool/types.rs | 10 + crates/sqlbuild-rules/src/lib.rs | 1 + python-boundary.toml | 31 +- src/sqlbuild/_native.pyi | 89 ++- .../_helpers/engine/custom_dispatch.py | 177 ----- .../_helpers/engine/custom_rules.py | 233 ------- .../_helpers/engine/rule_harness.py | 10 +- .../_helpers/host/custom_host_pool.py | 443 ------------- .../_helpers/host/host_capacity.py | 103 --- .../rule_engine/_helpers/run/rule_threads.py | 148 ----- .../rule_engine/classes/fact_digests.py | 52 -- src/sqlbuild/rule_engine/main/_evaluate.py | 44 -- src/sqlbuild/rule_engine/models.py | 32 +- src/sqlbuild/rule_engine/types.py | 5 - .../cli/commands/main/rules/helpers.py | 14 +- .../sqlbuild/cli/commands/main/test_rules.py | 94 +-- .../rule_engine/host_payload/_test_types.py | 8 - .../rule_engine/host_payload/helpers.py | 132 ---- .../host_payload/test_host_payload_facts.py | 136 ---- .../main/test_compile_phase_timings.py | 3 +- .../rule_engine/_helpers/engine/helpers.py | 9 +- .../_helpers/engine/test_catalogue.py | 2 +- .../engine/test_custom_rule_evidence.py | 2 +- .../_helpers/engine/test_custom_rule_reads.py | 13 +- .../_helpers/engine/test_custom_rules.py | 2 +- .../rule_engine/_helpers/host/helpers.py | 74 +-- .../_helpers/host/test_custom_host_pool.py | 397 +---------- .../_helpers/host/test_host_capacity.py | 134 ---- .../rule_engine/_helpers/run/helpers.py | 14 +- .../_helpers/run/test_native_rows.py | 8 +- .../evaluate/test_cache_and_suppressions.py | 2 +- .../main/evaluate/test_evaluate.py | 2 +- .../main/evaluate_rule/test_evaluate_rule.py | 2 +- 65 files changed, 2706 insertions(+), 2386 deletions(-) create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/rules/custom_evaluation.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/rules/rule_dispatch.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/_helpers/cgroups.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/_helpers/decoding.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/_helpers/merging.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/_helpers/mod.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/_helpers/payload_values.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/constants.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/main/available_cores.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/main/decode_host_response.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/main/host_count.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/main/merge_host_slices.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/main/mod.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/main/partition_host_plan.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/main/planned_invocations.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/main/without_codes.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/mod.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/models.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/tests.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/tests/helpers.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/tests/test_custom_host_pool.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/tests/test_types.rs create mode 100644 crates/sqlbuild-rules/src/custom_host_pool/types.rs delete mode 100644 src/sqlbuild/rule_engine/_helpers/engine/custom_dispatch.py delete mode 100644 src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py delete mode 100644 src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py delete mode 100644 src/sqlbuild/rule_engine/_helpers/host/host_capacity.py delete mode 100644 src/sqlbuild/rule_engine/_helpers/run/rule_threads.py delete mode 100644 src/sqlbuild/rule_engine/classes/fact_digests.py delete mode 100644 src/sqlbuild/rule_engine/main/_evaluate.py delete mode 100644 tests/integration/src/sqlbuild/rule_engine/host_payload/_test_types.py delete mode 100644 tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py delete mode 100644 tests/integration/src/sqlbuild/rule_engine/host_payload/test_host_payload_facts.py delete mode 100644 tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_host_capacity.py diff --git a/crates/sqlbuild-python/src/bindings/_helpers/cache/native_store.rs b/crates/sqlbuild-python/src/bindings/_helpers/cache/native_store.rs index 792448b1b7..63cd1b55b9 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/cache/native_store.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/cache/native_store.rs @@ -1,7 +1,7 @@ //! Digests, project fingerprints and store files of the shared native store, with the GIL released. use pyo3::exceptions::{PyOSError, PyRuntimeError}; -use pyo3::prelude::{Bound, PyAnyMethods, PyErr, PyModule, PyModuleMethods, PyResult, Python}; +use pyo3::prelude::{Bound, PyAnyMethods, PyModule, PyModuleMethods, PyResult, Python}; use pyo3::sync::PyOnceLock; use pyo3::{pyfunction, wrap_pyfunction}; use sqlbuild_cache::digest::main::compiled_code_identity::compiled_code_identity as compiled; @@ -14,7 +14,6 @@ use sqlbuild_cache::digest::main::installed_code_identity::installed_code_identi use sqlbuild_cache::digest::types::ContentDigest; use sqlbuild_cache::store::main::open_native_store::open_native_store; use sqlbuild_cache::store::models::NativeStore; -use sqlbuild_rules::constants::NATIVE_BUILD_IDENTITY; use std::collections::HashSet; use std::path::{Path, PathBuf}; @@ -91,12 +90,13 @@ pub(crate) fn save_store( /// The digest that changes with every released, locally edited, or rebuilt SQLBuild version. #[pyfunction] pub(crate) fn compiled_code_identity(py: Python<'_>) -> PyResult { - static IDENTITY: PyOnceLock = PyOnceLock::new(); - IDENTITY - .get_or_try_init(py, || { - Ok::(compiled(&installed_identity(py)?, NATIVE_BUILD_IDENTITY)) - }) - .cloned() + static INSTALLED: PyOnceLock = PyOnceLock::new(); + let installed: &String = INSTALLED.get_or_try_init(py, || installed_identity(py))?; + let build_identity: String = py + .import(NATIVE_MODULE)? + .getattr("BUILD_IDENTITY")? + .extract()?; + Ok(compiled(installed, &build_identity)) } /// `compiled_code_identity`, looked up on the extension module so tests can replace it. diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/analysis.rs b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/analysis.rs index 3cdcc843f3..f972159849 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/analysis.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/analysis.rs @@ -29,7 +29,7 @@ const FAST_COLUMN_LINEAGE: &str = "FAST"; const CONTRACTS_MODULE: &str = "sqlbuild._native"; const CONTRACT_MODELS_MODULE: &str = "sqlbuild.compiler.contracts.models"; const SELECTION_MODULE: &str = "sqlbuild._native"; -const RULE_THREADS_MODULE: &str = "sqlbuild.rule_engine._helpers.run.rule_threads"; +const RULE_THREADS_MODULE: &str = "sqlbuild._native"; const RULE_MODELS_MODULE: &str = "sqlbuild.rule_engine.models"; const COMPILE_MODELS_MODULE: &str = "sqlbuild.compiler.compile.models"; const COMPILE_TYPES_MODULE: &str = "sqlbuild.compiler.compile.types"; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/pipeline.rs b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/pipeline.rs index 8bf92a1d30..a53986c818 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/pipeline.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compile_driver/pipeline.rs @@ -26,7 +26,7 @@ const ANALYSIS_SELECTION_MODULE: &str = "sqlbuild._native"; const PIPELINE_MODELS_MODULE: &str = "sqlbuild.compiler.pipeline.models"; const TIMING_TRACKER_MODULE: &str = "sqlbuild._native"; const RULE_EXCEPTIONS_MODULE: &str = "sqlbuild.rule_engine.exceptions"; -const RULE_THREADS_MODULE: &str = "sqlbuild.rule_engine._helpers.run.rule_threads"; +const RULE_THREADS_MODULE: &str = "sqlbuild._native"; const RULE_MODELS_MODULE: &str = "sqlbuild.rule_engine.models"; const RUNTIME_MODELS_MODULE: &str = "sqlbuild.runtime.contracts.models"; const GENERIC_DIALECT: &str = "generic"; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/event_stream.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/event_stream.rs index bfeedb9399..911be9bf38 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/event_stream.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/event_stream.rs @@ -39,6 +39,10 @@ const STREAM_OPTIONS: (&[&str], &[&str]) = ( "enable_status", ], ); +/// A running node: unique id, display name and start offset in seconds. +type ActiveNode = (String, String, f64); +type LineResult = io::Result>; + const WAITING_MESSAGE: &str = "Waiting for dbt node output..."; /// How often the running-node progress line refreshes while dbt is silent. const STATUS_REFRESH: Duration = Duration::from_secs(1); @@ -50,7 +54,7 @@ const SIGNAL_POLL: Duration = Duration::from_millis(100); /// The live progress line and the running nodes it lists, in first-start order. struct StatusState { status: Option>, - active_nodes: Vec<(String, String, f64)>, + active_nodes: Vec, } /// What the read loop and the progress refresher share. @@ -173,14 +177,14 @@ fn streamed<'py>( let returncode: i64 = py .compiler_detach(|| Ok(dbt_process_returncode(&mut child))) .map_err(PyRuntimeError::new_err)? - .map_err(|error| PyErr::from(error))?; + .map_err(PyErr::from)?; waited = true; py.check_signals()?; Ok((returncode, results)) })(); - let child: Child = cleanup(py, &shared, refresher, (child, waited)); + let process_id: u32 = cleanup(py, &shared, refresher, (child, waited)); let (returncode, results) = result?; - Ok((returncode, results, Some(child.id()))) + Ok((returncode, results, Some(process_id))) } /// One dbt output line: progress message, node row, results and the running-node status. @@ -261,7 +265,7 @@ fn refreshed_status( py: Python<'_>, (stream, use_color, tty): (&Bound<'_, PyAny>, bool, bool), shared: &Shared, - (current, active_nodes): (Option>, &[(String, String, f64)]), + (current, active_nodes): (Option>, &[ActiveNode]), ) -> PyResult>> { if !tty || active_nodes.is_empty() { return Ok(None); @@ -306,10 +310,7 @@ fn locked<'a>( /// Read output lines on a thread so the waiting caller can notice interrupt signals. fn start_reader(mut lines: BufReader) -> Receiver>> { - let (sender, receiver): ( - Sender>>, - Receiver>>, - ) = channel(); + let (sender, receiver): (Sender, Receiver) = channel(); std::thread::spawn(move || { loop { let line: io::Result> = read_dbt_output_line(&mut lines); @@ -395,7 +396,7 @@ fn cleanup( shared: &Shared, refresher: Option<(JoinHandle<()>, Receiver<()>)>, (child, waited): (Child, bool), -) -> Child { +) -> u32 { if let Ok(mut stopped) = shared.stopped.lock() { *stopped = true; } @@ -415,9 +416,9 @@ fn cleanup( } let _closed: PyResult<()> = close_status(py, shared); if waited { - return child; + return child.id(); } - py.detach(|| stop_dbt_process(child, terminate)) + py.detach(|| stop_dbt_process(child, terminate)).id() } /// Ask a process to stop as `Popen.terminate` does. diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_evaluation.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_evaluation.rs new file mode 100644 index 0000000000..09a74008fe --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_evaluation.rs @@ -0,0 +1,495 @@ +//! Incremental custom-rule evaluation keyed on the compiler facts each invocation read. + +use std::collections::{BTreeMap, BTreeSet, HashMap}; +use std::path::PathBuf; +use std::time::Instant; + +use pyo3::IntoPyObjectExt; +use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyErr, PyResult, Python}; +use pyo3::types::{PyDict, PyDictMethods, PyTuple}; +use sqlbuild_cache::digest::main::sha256_hex::sha256_hex; +use sqlbuild_core::json::main::dumps::dumps; +use sqlbuild_core::json::models::{JsonDialect, JsonValue, OrjsonOptions}; +use sqlbuild_rules::custom_host_pool::models::{HostRuleSubject, HostRun}; +use sqlbuild_rules::custom_host_pool::types::{FactKey, HostPlan}; +use sqlbuild_rules::custom_rule_cache::models::CustomEvaluation; +use sqlbuild_rules::custom_rule_cache::types::ReadSet; + +use crate::bindings::_helpers::cache::native_store::current_code_identity; +use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported}; +use crate::bindings::_helpers::rules::custom_hosts::{HostRequest, run_custom_hosts}; +use crate::bindings::_helpers::rules::custom_rule_cache::CustomRuleCache; + +const NATIVE_MODULE: &str = "sqlbuild._native"; +const RULES_MODELS_MODULE: &str = "sqlbuild.rule_engine.models"; +const RULES_CONSTANTS_MODULE: &str = "sqlbuild.rule_engine.constants"; +const RULES_ERRORS_MODULE: &str = "sqlbuild.rule_engine.exceptions"; +const RULE_CONTEXT_MODULE: &str = "sqlbuild.rule_engine.classes.rule_context"; +const FACT_REPLAY_MODULE: &str = "sqlbuild.rule_engine._helpers.engine.fact_replay"; +const SCOPE_MODELS_MODULE: &str = "sqlbuild.compiler.scopes.models"; +const PROJECT_SUBJECT: &str = "project"; +const MODEL_SUBJECT: &str = "model"; +/// Project fields no rule fact reads, cleared before digests and host payloads use the project. +const HOST_PROJECT_CLEARED: [(&str, Cleared); 11] = [ + ("binding_catalog", Cleared::None), + ("sql_expansions", Cleared::Dict), + ("loaded_macros", Cleared::Dict), + ("loader_functions", Cleared::Tuple), + ("hook_functions", Cleared::Tuple), + ("materialization_files", Cleared::Tuple), + ("external_sql_reference_resolver", Cleared::None), + ("sql_scenarios", Cleared::Tuple), + ("sql_hook_files", Cleared::Tuple), + ("diagnostics", Cleared::Tuple), + ("unmatched_literal_sql_relations", Cleared::Tuple), +]; + +/// The empty value one cleared project field takes. +#[derive(Clone, Copy)] +enum Cleared { + None, + Dict, + Tuple, +} + +/// The run one custom-rule evaluation describes. +pub(crate) struct CustomRun<'a, 'py> { + pub(crate) project: &'a Bound<'py, PyAny>, + pub(crate) config: &'a Bound<'py, PyAny>, + pub(crate) project_dir: &'a Bound<'py, PyAny>, + pub(crate) rules: Vec>, + pub(crate) dialect: &'a str, + pub(crate) verify_determinism: bool, + pub(crate) implementation_fingerprints: HashMap, +} + +/// Reuse results whose recorded fact reads are unchanged and evaluate only the rest. +pub(crate) fn evaluate_custom_rules_cached<'py>( + run: CustomRun<'_, 'py>, +) -> PyResult> { + let py: Python<'py> = run.project.py(); + let started: Instant = Instant::now(); + let started_cpu: f64 = cpu_seconds(py)?; + let mut custom: Vec<(String, Bound<'py, PyAny>)> = run + .rules + .iter() + .map(|rule| Ok((rule.getattr("code")?.extract::()?, rule.clone()))) + .collect::>()?; + custom.sort_by(|left, right| left.0.cmp(&right.0)); + if custom.is_empty() { + return outcome(py, (Vec::new(), 0, 0), (0, 0)); + } + let facts_project: Bound<'py, PyAny> = host_project(run.project)?; + let enabled: bool = run + .config + .getattr("cache")? + .getattr("enabled")? + .is_truthy()?; + let mut rows: Vec<(String, Option, Option)> = Vec::new(); + for (code, rule) in &custom { + let identity: Option = if enabled { + Some(rule_identity( + &run, + rule, + run.implementation_fingerprints.get(code), + )?) + } else { + None + }; + rows.push((code.clone(), identity, subject_value(rule)?)); + } + let cache_path: Option = if enabled { + Some( + call_with( + &imported(py, NATIVE_MODULE, "rules_bulk_cache_path")?, + &[ + ("project_dir", run.project_dir.clone()), + ( + "file_name", + imported(py, RULES_CONSTANTS_MODULE, "CUSTOM_RULES_CACHE_FILE")?, + ), + ], + )? + .extract()?, + ) + } else { + None + }; + let mut cache: CustomRuleCache = CustomRuleCache::new( + cache_path, + rows, + run.project.getattr("effective_target_name")?.extract()?, + model_rows(run.project)?, + ); + let mut digests: FactDigests<'py> = + FactDigests::new(&facts_project, run.project_dir, run.dialect)?; + let pending: Vec = cache.pending_readsets().to_vec(); + let mut readset_digests: Vec> = Vec::new(); + for readset in &pending { + readset_digests.push(digests.reproducible(|digests| digests.combined(readset))?); + } + let plan: HostPlan = cache.plan(readset_digests); + if !plan.is_empty() { + let rules: Vec<(String, HostRuleSubject)> = custom + .iter() + .map(|(code, rule)| Ok((code.clone(), host_subject(rule)?))) + .collect::>()?; + let mut model_paths: Vec = model_rows(run.project)? + .into_iter() + .map(|(path, _, _)| path) + .collect(); + model_paths.sort(); + let project_dir: PathBuf = run.project_dir.extract()?; + let host_run: HostRun = run_custom_hosts( + py, + &HostRequest { + project: &facts_project, + config: run.config, + project_dir: &project_dir, + dialect: run.dialect, + rules: &rules, + model_paths: &model_paths, + track_reads: enabled, + verify_determinism: run.verify_determinism, + }, + &plan, + )?; + let evaluations: Vec = evaluations(host_run, &mut digests, enabled)?; + cache.record(py, evaluations)?; + } + let elapsed_ms: i64 = (started.elapsed().as_secs_f64() * 1000.0).round_ties_even() as i64; + let cpu_ms: i64 = ((cpu_seconds(py)? - started_cpu) * 1000.0) as i64; + outcome( + py, + (cache.findings(py)?, cache.hits(), cache.misses()), + (elapsed_ms, cpu_ms), + ) +} + +/// `CustomRulesOutcome(findings=..., cache_hits=..., cache_misses=..., custom_ms=..., ...)`. +pub(crate) fn outcome<'py>( + py: Python<'py>, + (findings, hits, misses): (Vec>, usize, usize), + (custom_ms, custom_cpu_ms): (i64, i64), +) -> PyResult> { + call_with( + &imported(py, RULES_MODELS_MODULE, "CustomRulesOutcome")?, + &[ + ("findings", PyTuple::new(py, findings)?.into_any()), + ("cache_hits", hits.into_bound_py_any(py)?), + ("cache_misses", misses.into_bound_py_any(py)?), + ("custom_ms", custom_ms.into_bound_py_any(py)?), + ("custom_cpu_ms", custom_cpu_ms.into_bound_py_any(py)?), + ], + ) +} + +/// Each evaluated invocation with its read set and its digest when it may be cached. +fn evaluations( + run: HostRun, + digests: &mut FactDigests<'_>, + cacheable: bool, +) -> PyResult> { + let mut full_keys: Option> = None; + let mut combined: HashMap> = HashMap::new(); + let mut rows: Vec = Vec::new(); + for evaluation in run.evaluated { + let mut stored: Option<(ReadSet, Option)> = None; + if cacheable && !run.uncacheable.contains(&evaluation.code) { + let reads: Vec = match evaluation.reads { + Some(reads) => reads, + None => { + if full_keys.is_none() { + full_keys = Some(digests.full_fact_keys()?); + } + let mut reads: Vec = run + .untracked_reads + .get(&evaluation.code) + .cloned() + .unwrap_or_default(); + reads.extend(full_keys.iter().flatten().cloned()); + reads + } + }; + let readset: ReadSet = reads + .into_iter() + .collect::>() + .into_iter() + .collect(); + if !combined.contains_key(&readset) { + let digest: Option = digests + .reproducible(|digests| digests.observed(&readset, &run.observed))? + .flatten(); + combined.insert(readset.clone(), digest); + } + let digest: Option = combined.get(&readset).cloned().flatten(); + stored = Some((readset, digest)); + } + rows.push(CustomEvaluation { + code: evaluation.code, + subject: evaluation.subject, + cacheable: stored, + findings: evaluation.findings, + }); + } + Ok(rows) +} + +/// Memoized digests of recorded fact keys against one compiled project snapshot. +struct FactDigests<'py> { + project: Bound<'py, PyAny>, + project_dir: Bound<'py, PyAny>, + dialect: String, + views: Option>, + shared: Bound<'py, PyAny>, + digests: HashMap, +} + +impl<'py> FactDigests<'py> { + fn new( + project: &Bound<'py, PyAny>, + project_dir: &Bound<'py, PyAny>, + dialect: &str, + ) -> PyResult { + Ok(Self { + project: project.clone(), + project_dir: project_dir.clone(), + dialect: dialect.to_owned(), + views: None, + shared: imported(project.py(), FACT_REPLAY_MODULE, "shared_fact_encoder")?.call0()?, + digests: HashMap::new(), + }) + } + + fn views(&mut self) -> PyResult> { + if let Some(views) = &self.views { + return Ok(views.clone()); + } + let py: Python<'py> = self.project.py(); + let views: Bound<'py, PyAny> = call_with( + &imported(py, RULE_CONTEXT_MODULE, "build_rule_fact_views")?, + &[ + ("project", self.project.clone()), + ("project_dir", self.project_dir.clone()), + ("dialect", self.dialect.clone().into_bound_py_any(py)?), + ], + )?; + self.views = Some(views.clone()); + Ok(views) + } + + /// One fact digest; raises `FactDigestError` when it cannot be reproduced. + fn digest(&mut self, key: &FactKey) -> PyResult { + if let Some(cached) = self.digests.get(key) { + return Ok(cached.clone()); + } + let py: Python<'py> = self.project.py(); + let views: Bound<'py, PyAny> = self.views()?; + let digest: String = call_with( + &imported(py, FACT_REPLAY_MODULE, "fact_outcome_digest")?, + &[ + ("views", views), + ("key", PyTuple::new(py, key)?.into_any()), + ("shared", self.shared.clone()), + ], + )? + .extract()?; + self.digests.insert(key.clone(), digest.clone()); + Ok(digest) + } + + /// One read set's digest over its keys in canonical order. + fn combined(&mut self, keys: &[FactKey]) -> PyResult { + let mut sorted: Vec<&FactKey> = keys.iter().collect(); + sorted.sort(); + let mut bytes: Vec = Vec::new(); + for key in sorted { + bytes.extend(orjson_key(key).into_bytes()); + bytes.extend(self.digest(key)?.into_bytes()); + } + Ok(sha256_hex(&bytes)) + } + + /// The read set's digest, or `None` when a filesystem fact changed since the rule read it. + fn observed( + &mut self, + readset: &[FactKey], + observed: &BTreeMap>, + ) -> PyResult> { + let py: Python<'py> = self.project.py(); + let is_tree_fact: Bound<'py, PyAny> = imported(py, FACT_REPLAY_MODULE, "is_tree_fact")?; + for key in readset { + if is_tree_fact.call1((PyTuple::new(py, key)?,))?.is_truthy()? + && let Some(seen) = observed.get(key) + && seen.as_deref() != Some(self.digest(key)?.as_str()) + { + return Ok(None); + } + } + self.combined(readset).map(Some) + } + + /// Keys covering every fact a custom rule can observe through its context. + fn full_fact_keys(&mut self) -> PyResult> { + let py: Python<'py> = self.project.py(); + imported(py, FACT_REPLAY_MODULE, "full_fact_keys")? + .call1((self.views()?,))? + .extract() + } + + /// `body`'s value, or `None` when it raised `FactDigestError`. + fn reproducible( + &mut self, + body: impl FnOnce(&mut Self) -> PyResult, + ) -> PyResult> { + let py: Python<'py> = self.project.py(); + match body(self) { + Ok(value) => Ok(Some(value)), + Err(error) if fact_digest_error(py, &error)? => Ok(None), + Err(error) => Err(error), + } + } +} + +fn fact_digest_error(py: Python<'_>, error: &PyErr) -> PyResult { + let class: Bound<'_, PyAny> = imported(py, RULES_ERRORS_MODULE, "FactDigestError")?; + error.value(py).is_instance(&class) +} + +/// `orjson.dumps(key)` of one fact key. +fn orjson_key(key: &FactKey) -> String { + let value: JsonValue = JsonValue::Array( + key.iter() + .map(|part| JsonValue::String(part.clone())) + .collect(), + ); + dumps(&value, &JsonDialect::Orjson(OrjsonOptions::default())).unwrap_or_default() +} + +/// The project without state no rule fact reads, before digests and host payloads use it. +fn host_project<'py>(project: &Bound<'py, PyAny>) -> PyResult> { + let py: Python<'py> = project.py(); + let options: Bound<'py, PyDict> = PyDict::new(py); + options.set_item("include_sql_expansions", false)?; + let record: Bound<'py, PyAny> = + imported(py, NATIVE_MODULE, "compiled_project_record")?.call((project,), Some(&options))?; + let changes: Bound<'py, PyDict> = PyDict::new(py); + for (name, cleared) in HOST_PROJECT_CLEARED { + match cleared { + Cleared::None => changes.set_item(name, py.None())?, + Cleared::Dict => changes.set_item(name, PyDict::new(py))?, + Cleared::Tuple => changes.set_item(name, PyTuple::empty(py))?, + } + } + changes.set_item( + "scope_index", + imported(py, SCOPE_MODELS_MODULE, "ScopeIndex")?.call0()?, + )?; + imported(py, "dataclasses", "replace")?.call((record,), Some(&changes)) +} + +/// `(relative path, name, materialization)` of each compiled model, in project order. +fn model_rows(project: &Bound<'_, PyAny>) -> PyResult)>> { + let mut rows: Vec<(String, String, Option)> = Vec::new(); + for model in project.getattr("models")?.try_iter()? { + let model: Bound<'_, PyAny> = model?; + let materialized: Bound<'_, PyAny> = model + .getattr("config")? + .getattr("values")? + .call_method1("get", ("materialized",))?; + rows.push(( + model + .getattr("relative_path")? + .call_method0("as_posix")? + .extract()?, + model.getattr("name")?.extract()?, + if materialized.is_instance_of::() { + Some(materialized.extract::()?) + } else { + None + }, + )); + } + Ok(rows) +} + +fn subject_value(rule: &Bound<'_, PyAny>) -> PyResult> { + let subject: Bound<'_, PyAny> = rule.getattr("subject")?; + if subject.is_none() { + return Ok(None); + } + subject.getattr("value")?.extract() +} + +fn host_subject(rule: &Bound<'_, PyAny>) -> PyResult { + Ok(match subject_value(rule)?.as_deref() { + Some(PROJECT_SUBJECT) => HostRuleSubject::Project, + Some(MODEL_SUBJECT) => HostRuleSubject::Model, + Some(_) | None => HostRuleSubject::Other, + }) +} + +fn rule_identity( + run: &CustomRun<'_, '_>, + rule: &Bound<'_, PyAny>, + fingerprint: Option<&String>, +) -> PyResult { + let py: Python<'_> = rule.py(); + let fingerprint: Bound<'_, PyAny> = match fingerprint { + Some(fingerprint) => fingerprint.into_bound_py_any(py)?, + None => { + let closure: Bound<'_, PyAny> = call_with( + &imported(py, NATIVE_MODULE, "custom_rule_import_closure")?, + &[ + ("rule", rule.clone()), + ("project_dir", run.project_dir.clone()), + ], + )?; + call_with( + &imported(py, NATIVE_MODULE, "custom_rule_implementation_fingerprint")?, + &[ + ("rule", rule.clone()), + ("project_dir", run.project_dir.clone()), + ("import_closure", closure), + ], + )? + } + }; + let subject: Option = subject_value(rule)?; + call_with( + &imported(py, NATIVE_MODULE, "custom_rule_identity")?, + &[ + ( + "code_identity", + current_code_identity(py)?.into_bound_py_any(py)?, + ), + ( + "runtime_version", + imported(py, RULES_CONSTANTS_MODULE, "CUSTOM_HOST_RUNTIME_VERSION")?, + ), + ("code", rule.getattr("code")?), + ("family", rule.getattr("family")?), + ("slug", rule.getattr("slug")?), + ("message", rule.getattr("message")?), + ("remediation", rule.getattr("remediation")?), + ("subject", subject.into_bound_py_any(py)?), + ("implementation_fingerprint", fingerprint), + ("config", run.config.clone()), + ("dialect", run.dialect.into_bound_py_any(py)?), + ( + "project_dir", + run.project_dir.call_method0("resolve")?.str()?.into_any(), + ), + ], + )? + .extract() +} + +/// This thread's CPU plus that of finished child processes, which run the rule hosts. +fn cpu_seconds(py: Python<'_>) -> PyResult { + let thread: f64 = imported(py, "time", "thread_time")?.call0()?.extract()?; + let times: Bound<'_, PyAny> = imported(py, "os", "times")?.call0()?; + let user: f64 = times.getattr("children_user")?.extract()?; + let system: f64 = times.getattr("children_system")?.extract()?; + Ok(thread + user + system) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs new file mode 100644 index 0000000000..a71aa3ee79 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs @@ -0,0 +1,399 @@ +//! Run planned custom-rule invocations across isolated host processes and merge their results. + +use std::collections::BTreeSet; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::mpsc::{Receiver, Sender, channel}; + +use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyErr, PyResult, Python}; +use pyo3::types::{PyBytes, PyBytesMethods, PyDict, PyDictMethods, PyTuple}; +use serde_json::{Map, Value, json}; +use sqlbuild_rules::custom_host_pool::constants::{CGROUP_ROOT, HOST_TIMEOUT_MILLIS, PROC_CGROUP}; +use sqlbuild_rules::custom_host_pool::main::available_cores::available_cores; +use sqlbuild_rules::custom_host_pool::main::decode_host_response::decode_host_response; +use sqlbuild_rules::custom_host_pool::main::host_count::host_count; +use sqlbuild_rules::custom_host_pool::main::merge_host_slices::merge_host_slices; +use sqlbuild_rules::custom_host_pool::main::partition_host_plan::partition_host_plan; +use sqlbuild_rules::custom_host_pool::main::planned_invocations::planned_invocations; +use sqlbuild_rules::custom_host_pool::main::without_codes::without_codes; +use sqlbuild_rules::custom_host_pool::models::{HostRuleSubject, HostRun, HostSlice}; +use sqlbuild_rules::custom_host_pool::types::{HostInvocation, HostPlan}; +use sqlbuild_rules::engine::main::custom_host::run_custom_host_json; + +use crate::bindings::_helpers::compile_driver::python_calls::imported; + +const NATIVE_MODULE: &str = "sqlbuild._native"; +const RULES_CONSTANTS_MODULE: &str = "sqlbuild.rule_engine.constants"; +const RULES_ERRORS_MODULE: &str = "sqlbuild.rule_engine.exceptions"; +const COMPILE_MODELS_MODULE: &str = "sqlbuild.compiler.compile.models"; +const HOST_INPUTS: [&str; 3] = ["target", "rules-cache", "host-inputs"]; +/// The model fields hosts decode; every other model field is dropped from their payload. +const PAYLOAD_MODEL_FIELDS: [&str; 18] = [ + "key", + "deps", + "name", + "relative_path", + "query_sql", + "config", + "destination", + "schema_entry", + "inferred_columns", + "fast_lineage_has_star", + "fast_lineage_star_resolved", + "authored_sql", + "extract_implicit_alias_columns", + "enum_declarations", + "constant_declarations", + "binding_validated", + "dynamic_column_contract", + "rejected_sql_analysis_opt_out", +]; + +static INPUT_COUNTER: AtomicUsize = AtomicUsize::new(0); + +/// What every host of one run is told, before its plan and cancellation marker are added. +pub(crate) struct HostRequest<'a, 'py> { + pub(crate) project: &'a Bound<'py, PyAny>, + pub(crate) config: &'a Bound<'py, PyAny>, + pub(crate) project_dir: &'a Path, + pub(crate) dialect: &'a str, + pub(crate) rules: &'a [(String, HostRuleSubject)], + pub(crate) model_paths: &'a [String], + pub(crate) track_reads: bool, + pub(crate) verify_determinism: bool, +} + +/// One host's outcome: its slice, or the `RulesError` message it failed with. +type HostOutcome = Result; +/// Whether a host finished, without its slice, so the waiter can cancel the rest. +type HostStatus = Result<(), String>; + +/// Evaluate `plan` across host processes and merge their results in single-host order. +pub(crate) fn run_custom_hosts( + py: Python<'_>, + request: &HostRequest<'_, '_>, + plan: &HostPlan, +) -> PyResult { + let invocations: Vec = + planned_invocations(plan, request.rules, request.model_paths); + let plans: Vec = if request.verify_determinism { + vec![plan.clone()] + } else { + partition_host_plan( + plan, + &invocations, + host_count(invocations.len(), cores(py)?), + ) + }; + let input_dir: PathBuf = HOST_INPUTS + .iter() + .fold(request.project_dir.to_path_buf(), |path, part| { + path.join(part) + }); + std::fs::create_dir_all(&input_dir)?; + let input_dir: PathBuf = input_dir.canonicalize()?; + let stem: String = format!("project-{}", uuid_hex(py)?); + let input_path: PathBuf = input_dir.join(format!("{stem}.pickle")); + let abandoned: PathBuf = input_dir.join(format!( + "{stem}{}", + constant::(py, "CUSTOM_HOST_ABANDONED_SUFFIX")? + )); + let mut payload: Map = Map::new(); + payload.insert("project_pickle_path".to_owned(), json!(text(&input_path))); + payload.insert("abandoned_marker".to_owned(), json!(text(&abandoned))); + payload.insert("parent_pid".to_owned(), json!(std::process::id())); + payload.insert( + "project_dir".to_owned(), + json!(text(&request.project_dir.canonicalize()?)), + ); + payload.insert("dialect".to_owned(), json!(request.dialect)); + payload.insert( + "verify_determinism".to_owned(), + json!(request.verify_determinism), + ); + payload.insert("track_reads".to_owned(), json!(request.track_reads)); + let launch: Launch = Launch::new(py, payload)?; + let responses: PyResult> = run_partitions( + py, + (&launch, plan, &plans), + &input_dir, + (|| write_inputs(py, &input_path, request), &abandoned), + ); + remove_file(&input_path); + remove_file(&abandoned); + Ok(merge_host_slices(responses?, &invocations)) +} + +/// The program, arguments and shared payload every host starts with. +struct Launch { + program: String, + arguments: Vec, + runtime_version: String, + payload: Map, +} + +impl Launch { + fn new(py: Python<'_>, payload: Map) -> PyResult { + Ok(Self { + program: imported(py, "sys", "executable")?.extract()?, + arguments: vec!["-m".to_owned(), constant(py, "CUSTOM_HOST_LAUNCH_MODULE")?], + runtime_version: constant(py, "CUSTOM_HOST_RUNTIME_VERSION")?, + payload, + }) + } + + /// The `run_custom_host_json` request for one host's plan and optional cancellation marker. + fn spec(&self, plan: &HostPlan, cancel_marker: Option<&Path>) -> String { + let mut payload: Map = self.payload.clone(); + payload.insert("plan".to_owned(), json!(plan)); + if let Some(marker) = cancel_marker { + payload.insert("cancel_marker".to_owned(), json!(text(marker))); + } + json!({ + "program": self.program, + "arguments": self.arguments, + "timeout_millis": HOST_TIMEOUT_MILLIS, + "runtime_version": self.runtime_version, + "payload": payload, + }) + .to_string() + } +} + +/// Start hosts while the payload is written; failures and module state fall back to one host. +fn run_partitions<'py>( + py: Python<'py>, + (launch, plan, plans): (&Launch, &HostPlan, &[HostPlan]), + input_dir: &Path, + (write_inputs, abandoned): (impl FnOnce() -> PyResult<()>, &Path), +) -> PyResult> { + if plans.len() == 1 { + let spec: String = launch.spec(plan, None); + let outcomes: Vec = + run_hosts(py, vec![spec], (write_inputs, abandoned), None)?; + return outcomes + .into_iter() + .map(|outcome| outcome.map_err(|message| rules_error(py, message))) + .collect(); + } + let cancel_dir: PathBuf = unique_directory(input_dir, "hosts-")?; + let marker: PathBuf = cancel_dir.join(constant::(py, "CUSTOM_HOST_CANCELLED_MARKER")?); + let specs: Vec = plans + .iter() + .map(|partition| launch.spec(partition, Some(&marker))) + .collect(); + let outcomes: PyResult> = + run_hosts(py, specs, (write_inputs, abandoned), Some(&marker)); + let _removed: std::io::Result<()> = std::fs::remove_dir_all(&cancel_dir); + let outcomes: Vec = outcomes?; + if outcomes.iter().any(Result::is_err) { + return single_host(py, launch, plan).map(|slice| vec![slice]); + } + let slices: Vec = outcomes.into_iter().filter_map(Result::ok).collect(); + let stateful: BTreeSet = slices + .iter() + .flat_map(|slice| slice.stateful.iter().cloned()) + .collect(); + if stateful.is_empty() { + return Ok(slices); + } + let rerun_plan: HostPlan = plan + .iter() + .map(|(code, subjects)| { + let subjects: Option> = if stateful.contains(code) { + subjects.clone() + } else { + Some(Vec::new()) + }; + (code.clone(), subjects) + }) + .collect(); + let rerun: HostSlice = single_host(py, launch, &rerun_plan)?; + let mut merged: Vec = slices + .into_iter() + .map(|slice| without_codes(slice, &stateful)) + .collect(); + merged.push(rerun); + Ok(merged) +} + +/// Run one host for `plan` once the payload already exists. +fn single_host(py: Python<'_>, launch: &Launch, plan: &HostPlan) -> PyResult { + let spec: String = launch.spec(plan, None); + py.detach(|| host_outcome(&spec)) + .map_err(|message| rules_error(py, message)) +} + +/// Run every host on its own thread while this thread publishes the payload they wait for. +fn run_hosts( + py: Python<'_>, + specs: Vec, + (write_inputs, abandoned): (impl FnOnce() -> PyResult<()>, &Path), + cancel_marker: Option<&Path>, +) -> PyResult> { + std::thread::scope(|scope| { + let (sender, finished): (Sender, Receiver) = channel(); + let handles: Vec<_> = specs + .into_iter() + .map(|spec| { + let sender: Sender = sender.clone(); + scope.spawn(move || { + let outcome: HostOutcome = host_outcome(&spec); + let status: HostStatus = outcome.as_ref().map(drop).map_err(Clone::clone); + let _sent: Result<(), std::sync::mpsc::SendError> = + sender.send(status); + outcome + }) + }) + .collect(); + drop(sender); + let published: PyResult<()> = write_inputs(); + if let Err(error) = published { + if let Ok(()) = touch(abandoned) { + py.detach(|| handles.into_iter().for_each(|handle| drop(handle.join()))); + } + return Err(error); + } + let count: usize = handles.len(); + py.detach(move || { + for _ in 0..count { + match finished.recv() { + Ok(Ok(())) => {} + Ok(Err(_)) | Err(_) => { + if let Some(marker) = cancel_marker { + let _touched: std::io::Result<()> = touch(marker); + } + break; + } + } + } + Ok(handles + .into_iter() + .map(|handle| { + handle + .join() + .unwrap_or_else(|_| Err("custom rule host thread panicked".to_owned())) + }) + .collect()) + }) + }) +} + +/// One host run decoded, or the message `RulesError` reports for it. +fn host_outcome(spec: &str) -> HostOutcome { + let response: String = run_custom_host_json(spec)?; + let value: Value = serde_json::from_str(&response).map_err(|error| error.to_string())?; + decode_host_response(&value) +} + +/// Publish the pickled project and config atomically; hosts started earlier wait for it. +fn write_inputs(py: Python<'_>, path: &Path, request: &HostRequest<'_, '_>) -> PyResult<()> { + let payload: Bound<'_, PyAny> = host_payload_project(request.project)?; + let pickled: Bound<'_, PyAny> = imported(py, "pickle", "dumps")? + .call1((PyTuple::new(py, [payload, request.config.clone()])?,))?; + let bytes: &[u8] = pickled.cast::()?.as_bytes(); + let parent: &Path = path.parent().unwrap_or(path); + let partial: PathBuf = parent.join(format!( + "writing-{}-{}.pickle", + std::process::id(), + INPUT_COUNTER.fetch_add(1, Ordering::Relaxed) + )); + if let Err(error) = std::fs::write(&partial, bytes) { + remove_file(&partial); + return Err(error.into()); + } + std::fs::rename(&partial, path)?; + Ok(()) +} + +/// The project hosts decode, without model state no rule fact reads. +fn host_payload_project<'py>(project: &Bound<'py, PyAny>) -> PyResult> { + let py: Python<'py> = project.py(); + let model_class: Bound<'py, PyAny> = imported(py, COMPILE_MODELS_MODULE, "CompiledModel")?; + let mut models: Vec> = Vec::new(); + for model in project.getattr("models")?.try_iter()? { + let model: Bound<'py, PyAny> = model?; + let fields: Bound<'py, PyDict> = PyDict::new(py); + for name in PAYLOAD_MODEL_FIELDS { + fields.set_item(name, model.getattr(name)?)?; + } + models.push(model_class.call((), Some(&fields))?); + } + let record: Bound<'py, PyAny> = + imported(py, NATIVE_MODULE, "compiled_project_record")?.call1((project,))?; + let options: Bound<'py, PyDict> = PyDict::new(py); + options.set_item("models", PyTuple::new(py, models)?)?; + options.set_item("native_project", py.None())?; + imported(py, "dataclasses", "replace")?.call((record,), Some(&options)) +} + +/// Usable CPUs: the process CPU count capped by any cgroup CPU quota. +fn cores(py: Python<'_>) -> PyResult { + let os: Bound<'_, PyAny> = py.import("os")?.into_any(); + let count: Option = if os.hasattr("process_cpu_count")? { + os.call_method0("process_cpu_count")?.extract()? + } else if os.hasattr("sched_getaffinity")? { + Some(os.call_method1("sched_getaffinity", (0,))?.len()?) + } else { + os.call_method0("cpu_count")?.extract()? + }; + Ok(available_cores( + count.unwrap_or(1), + (Path::new(CGROUP_ROOT), Path::new(PROC_CGROUP)), + )) +} + +fn constant<'py, T: pyo3::FromPyObject<'py>>(py: Python<'py>, name: &str) -> PyResult { + imported(py, RULES_CONSTANTS_MODULE, name)?.extract() +} + +fn uuid_hex(py: Python<'_>) -> PyResult { + imported(py, "uuid", "uuid4")? + .call0()? + .getattr("hex")? + .extract() +} + +/// A new directory below `parent` named `prefix` plus a unique suffix, as `mkdtemp` creates. +fn unique_directory(parent: &Path, prefix: &str) -> PyResult { + loop { + let directory: PathBuf = parent.join(format!( + "{prefix}{}-{}", + std::process::id(), + INPUT_COUNTER.fetch_add(1, Ordering::Relaxed) + )); + match std::fs::create_dir(&directory) { + Ok(()) => return Ok(directory), + Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {} + Err(error) => return Err(error.into()), + } + } +} + +/// Create `marker` and its folder so hosts can see it. +fn touch(marker: &Path) -> std::io::Result<()> { + if let Some(parent) = marker.parent() { + std::fs::create_dir_all(parent)?; + } + std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(marker) + .map(drop) +} + +fn remove_file(path: &Path) { + let _removed: std::io::Result<()> = std::fs::remove_file(path); +} + +fn text(path: &Path) -> String { + path.to_string_lossy().into_owned() +} + +fn rules_error(py: Python<'_>, message: String) -> PyErr { + match imported(py, RULES_ERRORS_MODULE, "RulesError").and_then(|class| class.call1((message,))) + { + Ok(instance) => PyErr::from_value(instance), + Err(error) => error, + } +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_rule_cache.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_rule_cache.rs index 8f8e0e3485..211128b211 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_rule_cache.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_rule_cache.rs @@ -4,8 +4,8 @@ use std::collections::BTreeMap; use std::path::PathBuf; use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyModule, PyModuleMethods, PyResult, Python}; -use pyo3::types::{PyDict, PyTuple}; -use pyo3::{pyclass, pyfunction, pymethods, wrap_pyfunction}; +use pyo3::types::PyDict; +use pyo3::{pyfunction, wrap_pyfunction}; use sqlbuild_rules::custom_rule_cache::main::custom_rule_cache_text::custom_rule_cache_text; use sqlbuild_rules::custom_rule_cache::main::lookup_custom_rule_cache::lookup_custom_rule_cache; use sqlbuild_rules::custom_rule_cache::main::ordered_custom_findings::ordered_custom_findings; @@ -28,14 +28,6 @@ use crate::bindings::_helpers::rules::findings::{FindingClasses, invalid_finding const RULES_ERRORS_MODULE: &str = "sqlbuild.rule_engine.exceptions"; const RULES_ERROR: &str = "RulesError"; -/// `(code, subject, (read set, digest) or None, findings)` of one evaluated invocation. -type EvaluationPy<'py> = ( - String, - String, - Option<(ReadSet, Option)>, - Vec>, -); - /// One custom rule's cache identity from its metadata, source fingerprint and settings. #[pyfunction] #[pyo3(signature = (**request))] @@ -57,14 +49,6 @@ fn custom_rule_identity(request: Option<&Bound<'_, PyDict>>) -> PyResult })) } -fn readset_tuple<'py>(py: Python<'py>, readset: &ReadSet) -> PyResult> { - let keys: Vec> = readset - .iter() - .map(|key| PyTuple::new(py, key)) - .collect::>()?; - PyTuple::new(py, keys) -} - fn rules_error(py: Python<'_>, message: String) -> PyResult { let class: Bound<'_, PyAny> = py.import(RULES_ERRORS_MODULE)?.getattr(RULES_ERROR)?; Ok(pyo3::PyErr::from_value(class.call1((message,))?)) @@ -77,9 +61,8 @@ fn stored_rules(path: &std::path::Path) -> StoredRules { parse_custom_rule_cache(&text) } -/// One custom-rule run's subjects, reused and new results, and the cache file it reads and writes. -#[pyclass(module = "sqlbuild._native")] -struct CustomRuleCache { +/// One run's custom-rule cache: stored results, pending read sets and accounting. +pub(crate) struct CustomRuleCache { path: Option, rules: Vec, project: SubjectsInput, @@ -90,12 +73,9 @@ struct CustomRuleCache { misses: usize, } -#[pymethods] impl CustomRuleCache { /// Load the cache at `path` (`None` when caching is off) for `(code, identity, subject)` rules. - #[new] - #[pyo3(signature = (*, path, rules, target_name, models))] - fn new( + pub(crate) fn new( path: Option, rules: Vec<(String, Option, Option)>, target_name: Option, @@ -137,16 +117,15 @@ impl CustomRuleCache { } /// Stored read sets whose current digest decides whether their results are reused. - fn pending_readsets<'py>(&self, py: Python<'py>) -> PyResult>> { - self.pending - .iter() - .map(|readset| readset_tuple(py, readset)) - .collect() + pub(crate) fn pending_readsets(&self) -> &[ReadSet] { + &self.pending } /// Reuse results given each pending read set's digest; returns the subjects still to run. - #[pyo3(signature = (*, digests))] - fn plan(&mut self, digests: Vec>) -> BTreeMap>> { + pub(crate) fn plan( + &mut self, + digests: Vec>, + ) -> BTreeMap>> { if self.path.is_none() { self.misses = self.rules.len(); return self @@ -170,19 +149,11 @@ impl CustomRuleCache { } /// Record evaluated results; with caching on, validate them, store them and write the file. - #[pyo3(signature = (*, evaluated))] - fn record(&mut self, py: Python<'_>, evaluated: Vec>) -> PyResult<()> { - let evaluations: Vec = evaluated - .into_iter() - .map(|(code, subject, cacheable, findings)| { - Ok(CustomEvaluation { - code, - subject, - cacheable, - findings: findings.iter().map(plain_value).collect::>()?, - }) - }) - .collect::>()?; + pub(crate) fn record( + &mut self, + py: Python<'_>, + evaluations: Vec, + ) -> PyResult<()> { let Some(path) = self.path.clone() else { self.extend_results(evaluations); return Ok(()); @@ -206,7 +177,7 @@ impl CustomRuleCache { } /// Every finding in rule code, then subject order; a malformed payload is a `RulesError`. - fn findings<'py>(&self, py: Python<'py>) -> PyResult>> { + pub(crate) fn findings<'py>(&self, py: Python<'py>) -> PyResult>> { let classes = FindingClasses::new(py)?; ordered_custom_findings(&self.rules, &self.results) .iter() @@ -218,18 +189,14 @@ impl CustomRuleCache { .collect() } - #[getter] - fn hits(&self) -> usize { + pub(crate) fn hits(&self) -> usize { self.hits } - #[getter] - fn misses(&self) -> usize { + pub(crate) fn misses(&self) -> usize { self.misses } -} -impl CustomRuleCache { fn extend_results(&mut self, evaluations: Vec) { for evaluation in evaluations { self.results @@ -239,7 +206,6 @@ impl CustomRuleCache { } pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { - module.add_class::()?; module.add_function(wrap_pyfunction!(custom_rule_identity, module)?)?; Ok(()) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/early_sql_lint.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/early_sql_lint.rs index 089cbdbf8f..46e1d43542 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/early_sql_lint.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/early_sql_lint.rs @@ -13,7 +13,7 @@ use sqlbuild_rules::rule_cache::main::rules_bulk_cache_path::rules_bulk_cache_pa use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported, keywords}; const NATIVE_MODULE: &str = "sqlbuild._native"; -const RULE_THREADS_MODULE: &str = "sqlbuild.rule_engine._helpers.run.rule_threads"; +const RULE_THREADS_MODULE: &str = "sqlbuild._native"; const RULES_MODELS_MODULE: &str = "sqlbuild.rule_engine.models"; const LINT_MODELS_MODULE: &str = "sqlbuild.lint.models"; const LINT_CONSTANTS_MODULE: &str = "sqlbuild.lint.constants"; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/mod.rs index f87e3ba6ed..2820c2e80c 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/mod.rs @@ -1,6 +1,8 @@ pub(crate) mod builtin_catalogue; pub(crate) mod catalogue; pub(crate) mod constants; +pub(crate) mod custom_evaluation; +pub(crate) mod custom_hosts; pub(crate) mod custom_rule_cache; pub(crate) mod custom_rule_evidence; pub(crate) mod dsl_headers; @@ -25,6 +27,7 @@ pub(crate) mod relation_keys; pub(crate) mod results; pub(crate) mod rows; pub(crate) mod rule_cache; +pub(crate) mod rule_dispatch; pub(crate) mod rule_plan; pub(crate) mod rule_run; pub(crate) mod rules_config; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/registration.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/registration.rs index adcb2564b4..ef27bfb097 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/registration.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/registration.rs @@ -63,5 +63,6 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { lint_stop::register(module)?; selection::register(module)?; results::register(module)?; - guidance::register(module) + guidance::register(module)?; + crate::bindings::_helpers::rules::rule_dispatch::register(module) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/rule_dispatch.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/rule_dispatch.rs new file mode 100644 index 0000000000..cacc14bd26 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/rule_dispatch.rs @@ -0,0 +1,621 @@ +//! Run built-in, SQL and custom rule phases on their threads and combine their findings. + +use std::collections::HashMap; +use std::thread::JoinHandle; + +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::{Bound, Py, PyAny, PyAnyMethods, PyModule, PyModuleMethods, PyResult, Python}; +use pyo3::types::{PyDict, PyDictMethods, PyFrozenSet, PyTuple}; +use pyo3::{IntoPyObjectExt, pyfunction, wrap_pyfunction}; + +use crate::bindings::_helpers::cache::native_store::current_code_identity; +use crate::bindings::_helpers::compile_driver::python_calls::{ + KeywordArguments, call_with, imported, within, +}; +use crate::bindings::_helpers::rules::custom_evaluation::{ + CustomRun, evaluate_custom_rules_cached, outcome, +}; + +const NATIVE_MODULE: &str = "sqlbuild._native"; +const RULES_MODELS_MODULE: &str = "sqlbuild.rule_engine.models"; +const RULES_CONSTANTS_MODULE: &str = "sqlbuild.rule_engine.constants"; +const CUSTOM_PAYLOADS_MODULE: &str = "sqlbuild.rule_engine._helpers.engine.custom_payloads"; +const HERMETICITY_MODULE: &str = "sqlbuild.rule_engine._helpers.engine.hermeticity"; +const CUSTOM_THREAD: &str = "sqlbuild-custom-rules"; +const GENERIC_DIALECT: &str = "generic"; +const PROJECT_RULES_OPTIONS: (&[&str], &[&str]) = ( + &["project", "config", "project_dir"], + &["dialect", "verify_determinism"], +); +const NATIVE_RULES_OPTIONS: (&[&str], &[&str]) = ( + &["project", "config", "project_dir", "catalogue"], + &[ + "dialect", + "initial_findings", + "defer_suppressions", + "verify_determinism", + ], +); +const CUSTOM_RULES_OPTIONS: (&[&str], &[&str]) = ( + &["project", "config", "project_dir", "rules", "dialect"], + &["verify_determinism", "implementation_fingerprints"], +); +const RULES_THREAD: &str = "sqlbuild-rules"; +const EVALUATE_RULES_OPTIONS: (&[&str], &[&str]) = ( + &[ + "graph", + "discovered_inputs", + "config", + "project_dir", + "dialect", + ], + &[ + "selected_keys", + "prepared_sql", + "expansion_reuse", + "no_sql_analysis", + ], +); +const PREPARE_OPTIONS: (&[&str], &[&str]) = ( + &[ + "project_dir", + "config", + "discovered_inputs", + "static_declaration_scope", + "compiled_expansions", + "stop", + ], + &[], +); + +/// A running custom-rule evaluation, whose result is a `CustomRulesOutcome`. +type CustomHandle = JoinHandle>>; + +/// The rule run one native evaluation describes; every field stays bound to its Python value. +struct NativeRun<'py> { + project: Bound<'py, PyAny>, + config: Bound<'py, PyAny>, + project_dir: Bound<'py, PyAny>, + catalogue: Bound<'py, PyAny>, + dialect: String, + verify_determinism: bool, + initial_findings: Option>, +} + +/// How a native evaluation obtains its custom-rule outcome. +enum CustomSource { + /// Already started (or nothing to start) by the caller. + Started(Option), + /// Start once the native request is built, beside the native evaluation. + Deferred, +} + +/// Evaluate the selected compiler rules over compiled models. +#[pyfunction] +#[pyo3(signature = (**options))] +fn evaluate_project_rules<'py>( + py: Python<'py>, + options: Option>, +) -> PyResult> { + let arguments: KeywordArguments<'py> = + KeywordArguments::read(py, "evaluate_project_rules", options, PROJECT_RULES_OPTIONS)?; + let project: Bound<'py, PyAny> = arguments.value("project")?; + let config: Bound<'py, PyAny> = arguments.value("config")?; + let project_dir: Bound<'py, PyAny> = arguments.value("project_dir")?; + let dialect: String = dialect_argument(&arguments)?; + let verify_determinism: bool = arguments.flag("verify_determinism")?; + let effective: Bound<'py, PyAny> = call_with( + &native(py, "resolve_rule_ignore_selectors")?, + &[("config", config), ("project", project.clone())], + )?; + let catalogue: Bound<'py, PyAny> = call_with( + &native(py, "build_rule_catalogue")?, + &[ + ("config", effective.clone()), + ("project_dir", project_dir.clone()), + ], + )?; + let mut custom: bool = false; + for rule in catalogue.try_iter()? { + custom = custom || rule?.getattr("custom")?.is_truthy()?; + } + if custom { + let rules: Bound<'py, PyAny> = call_with( + &native(py, "select_catalogue_rules")?, + &[ + ("catalogue", catalogue.clone()), + ("config", effective.clone()), + ("project_dir", project_dir.clone()), + ], + )?; + call_with( + &imported(py, HERMETICITY_MODULE, "verify_custom_rules")?, + &[("rules", rules), ("project_dir", project_dir.clone())], + )?; + } + let run: NativeRun<'py> = NativeRun { + project, + config: effective, + project_dir, + catalogue, + dialect, + verify_determinism, + initial_findings: None, + }; + evaluate_native(&run, (false, None), CustomSource::Deferred) +} + +/// Evaluate built-in rules natively with the catalogue's selected custom rules beside them. +#[pyfunction] +#[pyo3(signature = (**options))] +fn evaluate_native_rules<'py>( + py: Python<'py>, + options: Option>, +) -> PyResult> { + let arguments: KeywordArguments<'py> = + KeywordArguments::read(py, "evaluate_native_rules", options, NATIVE_RULES_OPTIONS)?; + let initial_findings: Bound<'py, PyAny> = arguments.value("initial_findings")?; + let defer_suppressions: bool = arguments.flag("defer_suppressions")?; + let run: NativeRun<'py> = NativeRun { + project: arguments.value("project")?, + config: arguments.value("config")?, + project_dir: arguments.value("project_dir")?, + catalogue: arguments.value("catalogue")?, + dialect: dialect_argument(&arguments)?, + verify_determinism: arguments.flag("verify_determinism")?, + initial_findings: (!initial_findings.is_none()).then_some(initial_findings), + }; + evaluate_native(&run, (defer_suppressions, None), CustomSource::Deferred) +} + +/// Evaluate custom rules through their incremental cache. +#[pyfunction] +#[pyo3(signature = (**options))] +fn evaluate_custom_rules<'py>( + py: Python<'py>, + options: Option>, +) -> PyResult> { + let arguments: KeywordArguments<'py> = + KeywordArguments::read(py, "evaluate_custom_rules", options, CUSTOM_RULES_OPTIONS)?; + let fingerprints: Bound<'py, PyAny> = arguments.value("implementation_fingerprints")?; + evaluate_custom_rules_cached(CustomRun { + project: &arguments.value("project")?, + config: &arguments.value("config")?, + project_dir: &arguments.value("project_dir")?, + rules: arguments + .value("rules")? + .try_iter()? + .collect::>()?, + dialect: &arguments.value("dialect")?.extract::()?, + verify_determinism: arguments.flag("verify_determinism")?, + implementation_fingerprints: if fingerprints.is_none() { + HashMap::new() + } else { + fingerprints.extract()? + }, + }) +} + +/// The `dialect` argument, `"generic"` when it was not passed. +fn dialect_argument(arguments: &KeywordArguments<'_>) -> PyResult { + if arguments.passed("dialect")? { + return arguments.value("dialect")?.extract(); + } + Ok(GENERIC_DIALECT.to_owned()) +} + +/// Evaluate independent rule phases concurrently, then finalize their combined findings. +#[pyfunction] +#[pyo3(signature = (**options))] +pub(crate) fn evaluate_rules<'py>( + py: Python<'py>, + options: Option>, +) -> PyResult> { + let arguments: KeywordArguments<'py> = + KeywordArguments::read(py, "evaluate_rules", options, EVALUATE_RULES_OPTIONS)?; + let graph: Bound<'py, PyAny> = arguments.value("graph")?; + let dialect: Bound<'py, PyAny> = arguments.value("dialect")?; + let plan: Bound<'py, PyAny> = call_with( + &native(py, "plan_rule_run")?, + &[ + ("graph", graph.clone()), + ("config", arguments.value("config")?), + ("project_dir", arguments.value("project_dir")?), + ("selected_keys", arguments.value("selected_keys")?), + ( + "no_sql_analysis", + arguments.value_or_false("no_sql_analysis")?, + ), + ], + )?; + let run: NativeRun<'py> = NativeRun { + project: plan.getattr("project")?, + config: plan.getattr("config")?, + project_dir: plan.getattr("project_dir")?, + catalogue: plan.getattr("catalogue")?, + dialect: dialect.extract()?, + verify_determinism: false, + initial_findings: None, + }; + let payloads: Bound<'py, PyAny> = plan.getattr("custom_payloads")?; + let custom: Option = start_custom_rules(&run, &payloads)?; + let native_run: [Py; 5] = [ + run.project.clone().unbind(), + run.config.clone().unbind(), + run.project_dir.clone().unbind(), + run.catalogue.clone().unbind(), + payloads.clone().unbind(), + ]; + let native_dialect: String = run.dialect.clone(); + let native_thread: JoinHandle>> = std::thread::Builder::new() + .name(RULES_THREAD.to_owned()) + .spawn(move || { + Python::attach(|py| { + let [project, config, project_dir, catalogue, payloads] = native_run; + let run: NativeRun<'_> = NativeRun { + project: project.into_bound(py), + config: config.into_bound(py), + project_dir: project_dir.into_bound(py), + catalogue: catalogue.into_bound(py), + dialect: native_dialect, + verify_determinism: false, + initial_findings: None, + }; + evaluate_native( + &run, + (true, Some(payloads.into_bound(py))), + CustomSource::Started(custom), + ) + .map(Bound::unbind) + }) + }) + .map_err(|error| PyRuntimeError::new_err(error.to_string()))?; + let sql_result: PyResult> = call_with( + &native(py, "run_planned_sql_rules")?, + &[ + ("plan", plan.clone()), + ("graph", graph), + ("discovered_inputs", arguments.value("discovered_inputs")?), + ("dialect", dialect), + ("prepared_sql", arguments.value("prepared_sql")?), + ("expansion_reuse", arguments.value("expansion_reuse")?), + ], + ); + let result: PyResult> = joined(py, native_thread); + let sql_result: Bound<'py, PyAny> = sql_result?; + let result: Bound<'py, PyAny> = result?; + call_with( + &native(py, "finish_rule_run")?, + &[ + ("plan", plan), + ("sql_result", sql_result), + ("result", result), + ], + ) +} + +/// Start early SQL lint on `executor` when the project selects SQL rules without a cache. +#[pyfunction] +#[pyo3(signature = (*, inputs, executor, dialect, stop))] +pub(crate) fn prepare_sql_rules<'py>( + inputs: &Bound<'py, PyAny>, + executor: &Bound<'py, PyAny>, + dialect: &Bound<'py, PyAny>, + stop: &Bound<'py, PyAny>, +) -> PyResult> { + let py: Python<'py> = inputs.py(); + let preparation: Bound<'py, PyAny> = call_with( + &native(py, "sql_lint_preparation")?, + &[("inputs", inputs.clone()), ("dialect", dialect.clone())], + )?; + if preparation.is_none() { + return Ok(preparation); + } + let (project_dir, codes, config): (Bound<'py, PyAny>, Bound<'py, PyAny>, Bound<'py, PyAny>) = + preparation.extract()?; + let expansions: Bound<'py, PyAny> = + native(py, "early_lint_expansions")?.call1((inputs.getattr("model_inputs")?,))?; + let work: Bound<'py, PyDict> = PyDict::new(py); + work.set_item("project_dir", &project_dir)?; + work.set_item("config", config)?; + work.set_item("discovered_inputs", inputs.getattr("discovered_inputs")?)?; + work.set_item( + "static_declaration_scope", + inputs.getattr("declaration_scope")?, + )?; + work.set_item("compiled_expansions", expansions)?; + work.set_item("stop", stop)?; + let future: Bound<'py, PyAny> = executor.call_method( + "submit", + (wrap_pyfunction!(prepare_early_sql_lint, py)?,), + Some(&work), + )?; + call_with( + &imported(py, RULES_MODELS_MODULE, "PreparedSqlLint")?, + &[ + ("codes", codes), + ("project_dir", project_dir), + ("future", future), + ], + ) +} + +/// Expand and lint SQL early, stopping when `stop` is set; runs on the early-lint executor. +#[pyfunction] +#[pyo3(signature = (**options))] +fn prepare_early_sql_lint<'py>( + py: Python<'py>, + options: Option>, +) -> PyResult> { + let arguments: KeywordArguments<'py> = + KeywordArguments::read(py, "prepare_early_sql_lint", options, PREPARE_OPTIONS)?; + let project_dir: Bound<'py, PyAny> = arguments.value("project_dir")?; + let discovered_inputs: Bound<'py, PyAny> = arguments.value("discovered_inputs")?; + let scope: Bound<'py, PyAny> = + native(py, "lint_stop_scope")?.call1((arguments.value("stop")?,))?; + within(&scope, |_| { + let context: Bound<'py, PyAny> = call_with( + &native(py, "build_lint_expansion_context")?, + &[ + ("project_dir", project_dir.clone()), + ("discovered_inputs", discovered_inputs.clone()), + ( + "static_declaration_scope", + arguments.value("static_declaration_scope")?, + ), + ], + )?; + let result: Bound<'py, PyAny> = call_with( + &native(py, "run_lint")?, + &[ + ("project_dir", project_dir.clone()), + ("config", arguments.value("config")?), + ("discovered_inputs", discovered_inputs.clone()), + ( + "compiled_expansions", + arguments.value("compiled_expansions")?, + ), + ("expansion_context", context.clone()), + ], + )?; + call_with( + &imported(py, RULES_MODELS_MODULE, "PreparedSqlLintResult")?, + &[("result", result), ("context", context)], + ) + }) +} + +/// Evaluate built-in rules natively while custom rules run beside them. +fn evaluate_native<'py>( + run: &NativeRun<'py>, + (defer_suppressions, custom_payloads): (bool, Option>), + custom: CustomSource, +) -> PyResult> { + let py: Python<'py> = run.project.py(); + let prepared: PyResult<(Bound<'py, PyAny>, Bound<'py, PyAny>)> = + native_request(run, custom_payloads); + let custom: Option = match (custom, &prepared) { + (CustomSource::Started(handle), Ok(_)) => handle, + (CustomSource::Started(handle), Err(_)) => { + if let Some(handle) = handle { + let _abandoned: PyResult> = joined(py, handle); + } + None + } + (CustomSource::Deferred, Ok((payloads, _))) => start_custom_rules(run, payloads)?, + (CustomSource::Deferred, Err(_)) => None, + }; + let (_, request): (Bound<'py, PyAny>, Bound<'py, PyAny>) = prepared?; + let evaluated: PyResult> = + native(py, "evaluate_rules_request").and_then(|evaluate| evaluate.call1((request,))); + let custom: PyResult> = match custom { + Some(handle) => joined(py, handle), + None => outcome(py, (Vec::new(), 0, 0), (0, 0)), + }; + let native_result: Bound<'py, PyAny> = evaluated?; + let custom: Bound<'py, PyAny> = custom?; + let mut findings: Bound<'py, PyAny> = PyTuple::new( + py, + native_result + .getattr("findings")? + .try_iter()? + .chain(custom.getattr("findings")?.try_iter()?) + .collect::>>()?, + )? + .into_any(); + if !defer_suppressions { + findings = call_with( + &native(py, "finalize_rule_findings_rows")?, + &[ + ("config", run.config.clone()), + ("project_dir", run.project_dir.clone()), + ("evaluated_codes", native_result.getattr("selected_codes")?), + ("findings", findings), + ], + )?; + } + let reused: bool = native_result.getattr("reused")?.is_truthy()?; + let native_hits: i64 = native_result.getattr("cache_hits")?.extract()?; + let native_misses: i64 = native_result.getattr("cache_misses")?.extract()?; + let custom_hits: i64 = custom.getattr("cache_hits")?.extract()?; + let custom_misses: i64 = custom.getattr("cache_misses")?.extract()?; + call_with( + &imported(py, RULES_MODELS_MODULE, "RulesResult")?, + &[ + ("findings", findings), + ( + "evaluated_models", + native_result.getattr("evaluated_models")?, + ), + ( + "cache_hits", + (if reused { + native_hits + native_misses + } else { + native_hits + } + custom_hits) + .into_bound_py_any(py)?, + ), + ( + "cache_misses", + (if reused { 0 } else { native_misses } + custom_misses).into_bound_py_any(py)?, + ), + ( + "built_in_ms", + if reused { + 0.into_bound_py_any(py)? + } else { + native_result.getattr("built_in_ms")? + }, + ), + ("custom_ms", custom.getattr("custom_ms")?), + ("custom_cpu_ms", custom.getattr("custom_cpu_ms")?), + ], + ) +} + +/// The custom-rule payloads and the native rules request of one evaluation. +fn native_request<'py>( + run: &NativeRun<'py>, + custom_payloads: Option>, +) -> PyResult<(Bound<'py, PyAny>, Bound<'py, PyAny>)> { + let py: Python<'py> = run.project.py(); + let payloads: Bound<'py, PyAny> = match custom_payloads { + Some(payloads) => payloads, + None => call_with( + &imported(py, CUSTOM_PAYLOADS_MODULE, "custom_rule_payloads")?, + &[ + ("catalogue", run.catalogue.clone()), + ("project_dir", run.project_dir.clone()), + ], + )?, + }; + let mut include_type_proof: bool = false; + for code in imported(py, RULES_CONSTANTS_MODULE, "TYPE_PROOF_RULE_CODES")?.try_iter()? { + include_type_proof = include_type_proof + || call_with( + &native(py, "rule_selected")?, + &[("config", run.config.clone()), ("code", code?)], + )? + .is_truthy()?; + } + let request: Bound<'py, PyAny> = call_with( + &native(py, "build_rules_request")?, + &[ + ("project", run.project.clone()), + ("config", run.config.clone()), + ("project_dir", run.project_dir.clone()), + ("dialect", run.dialect.clone().into_bound_py_any(py)?), + ( + "initial_findings", + run.initial_findings + .clone() + .unwrap_or_else(|| PyTuple::empty(py).into_any()), + ), + ("custom_payloads", payloads.clone()), + ( + "include_type_proof", + include_type_proof.into_bound_py_any(py)?, + ), + ], + )?; + Ok((payloads, request)) +} + +/// Start the selected custom rules on their thread after digesting installed code here. +fn start_custom_rules( + run: &NativeRun<'_>, + payloads: &Bound<'_, PyAny>, +) -> PyResult> { + let py: Python<'_> = run.project.py(); + if !payloads.is_truthy()? { + return Ok(None); + } + let selected: Bound<'_, PyFrozenSet> = PyFrozenSet::new( + py, + call_with( + &native(py, "selected_rule_codes")?, + &[ + ("config", run.config.clone()), + ("custom_payloads", payloads.clone()), + ], + )? + .try_iter()? + .collect::>>()?, + )?; + let mut rules: Vec> = Vec::new(); + for rule in run.catalogue.try_iter()? { + let rule: Bound<'_, PyAny> = rule?; + if rule.getattr("custom")?.is_truthy()? && selected.contains(rule.getattr("code")?)? { + rules.push(rule.unbind()); + } + } + if rules.is_empty() { + return Ok(None); + } + if run + .config + .getattr("cache")? + .getattr("enabled")? + .is_truthy()? + { + let _identity: String = current_code_identity(py)?; + } + let mut fingerprints: HashMap = HashMap::new(); + for payload in payloads.try_iter()? { + let payload: Bound<'_, PyAny> = payload?; + fingerprints.insert( + payload.get_item("code")?.str()?.extract()?, + payload + .get_item("implementation_fingerprint")? + .str()? + .extract()?, + ); + } + let inputs: [Py; 3] = [ + run.project.clone().unbind(), + run.config.clone().unbind(), + run.project_dir.call_method0("resolve")?.unbind(), + ]; + let dialect: String = run.dialect.clone(); + let verify_determinism: bool = run.verify_determinism; + std::thread::Builder::new() + .name(CUSTOM_THREAD.to_owned()) + .spawn(move || { + Python::attach(|py| { + let [project, config, project_dir] = inputs; + evaluate_custom_rules_cached(CustomRun { + project: project.bind(py), + config: config.bind(py), + project_dir: project_dir.bind(py), + rules: rules.into_iter().map(|rule| rule.into_bound(py)).collect(), + dialect: &dialect, + verify_determinism, + implementation_fingerprints: fingerprints, + }) + .map(Bound::unbind) + }) + }) + .map(Some) + .map_err(|error| PyRuntimeError::new_err(error.to_string())) +} + +/// Wait for a phase thread without holding the interpreter, then re-raise what it raised. +fn joined<'py>(py: Python<'py>, handle: CustomHandle) -> PyResult> { + py.detach(|| handle.join()) + .map_err(|_| PyRuntimeError::new_err("rules thread panicked"))? + .map(|value| value.into_bound(py)) +} + +fn native<'py>(py: Python<'py>, name: &str) -> PyResult> { + imported(py, NATIVE_MODULE, name) +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(evaluate_project_rules, module)?)?; + module.add_function(wrap_pyfunction!(evaluate_native_rules, module)?)?; + module.add_function(wrap_pyfunction!(evaluate_custom_rules, module)?)?; + module.add_function(wrap_pyfunction!(evaluate_rules, module)?)?; + module.add_function(wrap_pyfunction!(prepare_sql_rules, module)?)?; + module.add_function(wrap_pyfunction!(prepare_early_sql_lint, module)?) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/rule_run.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/rule_run.rs index d9406540f6..cb2d1f2acd 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/rule_run.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/rule_run.rs @@ -15,7 +15,6 @@ use crate::bindings::_helpers::compile_driver::python_calls::{call_with, importe use crate::bindings::_helpers::rules::results::rule_finding; use crate::bindings::_helpers::rules::rows::finalize_rule_findings_rows; -const RULE_THREADS_MODULE: &str = "sqlbuild.rule_engine._helpers.run.rule_threads"; const MODELS_MODULE: &str = "sqlbuild.rule_engine.models"; const PATHLIB_MODULE: &str = "pathlib"; const CTYPES_MODULE: &str = "ctypes"; @@ -132,7 +131,7 @@ pub(crate) fn run_rules<'py>( py: Python<'py>, options: Option>, ) -> PyResult> { - imported(py, RULE_THREADS_MODULE, "evaluate_rules")?.call((), options.as_ref()) + crate::bindings::_helpers::rules::rule_dispatch::evaluate_rules(py, options) } /// One bulk Rules cache file owned by the native compiler engine. diff --git a/crates/sqlbuild-rules/src/custom_host_pool/_helpers/cgroups.rs b/crates/sqlbuild-rules/src/custom_host_pool/_helpers/cgroups.rs new file mode 100644 index 0000000000..81440f52b9 --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/_helpers/cgroups.rs @@ -0,0 +1,83 @@ +//! cgroup v1 and v2 CPU quotas along this process's cgroup paths. + +use std::path::{Path, PathBuf}; + +use crate::custom_host_pool::constants::{CGROUP_V1_CPU_CONTROLLER, CGROUP_V1_CPU_DIRECTORIES}; + +/// Every CPU quota, in whole CPUs rounded up, along each cgroup membership listed in `entries`. +pub(crate) fn cgroup_quotas(cgroup_root: &Path, entries: &[String]) -> Vec { + let fallback: [String; 1] = ["0::/".to_owned()]; + let entries: &[String] = if entries.is_empty() { + &fallback + } else { + entries + }; + let mut quotas: Vec = Vec::new(); + for entry in entries { + let remainder: &str = entry.split_once(':').map_or("", |(_, rest)| rest); + let (controllers, path): (&str, &str) = + remainder.split_once(':').unwrap_or((remainder, "")); + if controllers.is_empty() { + for directory in directories(cgroup_root, path) { + quotas.extend(cpu_max_quota(&directory.join("cpu.max"))); + } + } else if controllers + .split(',') + .any(|name| name == CGROUP_V1_CPU_CONTROLLER) + { + for name in std::iter::once(controllers).chain(CGROUP_V1_CPU_DIRECTORIES) { + for directory in directories(&cgroup_root.join(name), path) { + quotas.extend(cfs_quota(&directory)); + } + } + } + } + quotas +} + +/// The cgroup directory's ancestors from the hierarchy root down to it. +fn directories(base: &Path, relative: &str) -> Vec { + let mut directory: PathBuf = base.to_path_buf(); + let mut found: Vec = vec![directory.clone()]; + for part in Path::new(relative.trim_matches('/')).components() { + directory = directory.join(part); + found.push(directory.clone()); + } + found +} + +/// cgroup v2 `cpu.max`: `max ` is unlimited, else ` `. +fn cpu_max_quota(path: &Path) -> Option { + let Ok(text) = std::fs::read_to_string(path) else { + return None; + }; + let fields: Vec<&str> = text.split_whitespace().collect(); + match fields.as_slice() { + [quota, period] => quota_cpus(quota, period), + _ => None, + } +} + +/// cgroup v1 `cpu.cfs_quota_us` and `cpu.cfs_period_us`; a negative quota is unlimited. +fn cfs_quota(directory: &Path) -> Option { + let (Ok(quota), Ok(period)) = ( + std::fs::read_to_string(directory.join("cpu.cfs_quota_us")), + std::fs::read_to_string(directory.join("cpu.cfs_period_us")), + ) else { + return None; + }; + quota_cpus(quota.trim(), period.trim()) +} + +fn quota_cpus(quota: &str, period: &str) -> Option { + let (Ok(quota), Ok(period)) = (quota.parse::(), period.parse::()) else { + return None; + }; + if quota <= 0 || period <= 0 { + return None; + } + match usize::try_from(quota.div_euclid(period) + i64::from(quota.rem_euclid(period) != 0)) { + Ok(cpus) => Some(cpus.max(1)), + Err(_) => None, + } +} diff --git a/crates/sqlbuild-rules/src/custom_host_pool/_helpers/decoding.rs b/crates/sqlbuild-rules/src/custom_host_pool/_helpers/decoding.rs new file mode 100644 index 0000000000..a6f3b24df4 --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/_helpers/decoding.rs @@ -0,0 +1,119 @@ +//! Host response decoding. + +use std::collections::{BTreeMap, BTreeSet}; + +use serde_json::{Map, Value}; + +use crate::custom_host_pool::_helpers::payload_values::{fact_key, python_text, readset}; +use crate::custom_host_pool::constants::{ + CUSTOM_RULE_PROJECT_SUBJECT, INVALID_EVALUATION, INVALID_FINDINGS, INVALID_PAYLOAD, +}; +use crate::custom_host_pool::models::{HostEvaluation, HostSlice}; +use crate::custom_host_pool::types::FactKey; + +/// The slice a host response describes, or the `RulesError` message of a malformed one. +pub(crate) fn decoded_slice(response: &Value) -> Result { + let Value::Object(payload) = response else { + return Err(INVALID_PAYLOAD.to_owned()); + }; + let ( + Some(Value::Array(raw_readsets)), + Some(Value::Array(raw_evaluations)), + Some(Value::Object(raw_untracked)), + Some(Value::Array(raw_uncacheable)), + Some(Value::Array(raw_observed)), + ) = ( + payload.get("readsets"), + payload.get("evaluations"), + payload.get("untracked"), + payload.get("uncacheable"), + payload.get("observed"), + ) + else { + return Err(INVALID_PAYLOAD.to_owned()); + }; + let readsets: Vec> = raw_readsets.iter().map(readset).collect(); + let untracked_reads: BTreeMap> = raw_untracked + .iter() + .map(|(code, reads)| (code.clone(), readset(reads))) + .collect(); + let evaluations: Vec = raw_evaluations + .iter() + .map(|item| evaluation(item, &readsets, &untracked_reads)) + .collect::>()?; + Ok(HostSlice { + evaluations, + uncacheable: raw_uncacheable.iter().map(python_text).collect(), + observed: observed(raw_observed), + bounded: payload.get("bounded") == Some(&Value::Bool(true)), + stateful: match payload.get("stateful") { + Some(Value::Array(codes)) => codes.iter().map(python_text).collect(), + Some(_) | None => BTreeSet::new(), + }, + untracked_reads, + }) +} + +fn evaluation( + item: &Value, + readsets: &[Vec], + untracked: &BTreeMap>, +) -> Result { + let Value::Object(item) = item else { + return Err(INVALID_EVALUATION.to_owned()); + }; + let code: String = item + .get("code") + .map(python_text) + .ok_or(INVALID_EVALUATION)?; + let Some(Value::Array(findings)) = item.get("findings") else { + return Err(INVALID_FINDINGS.to_owned()); + }; + let reads: Option> = match item.get("reads") { + Some(Value::Number(index)) if !untracked.contains_key(&code) => match index.as_u64() { + Some(index) => usize::try_from(index) + .map_or(None, |index| readsets.get(index)) + .cloned(), + None => None, + }, + Some(_) | None => None, + }; + Ok(HostEvaluation { + subject: match item.get("subject") { + None | Some(Value::Null) => CUSTOM_RULE_PROJECT_SUBJECT.to_owned(), + Some(subject) => python_text(subject), + }, + findings: findings.iter().map(finding).collect(), + reads, + code, + }) +} + +fn finding(value: &Value) -> Value { + match value { + Value::Object(fields) => Value::Object(fields.clone()), + Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) | Value::Array(_) => { + Value::Object(Map::new()) + } + } +} + +fn observed(raw: &[Value]) -> BTreeMap> { + let mut observed: BTreeMap> = BTreeMap::new(); + for entry in raw { + if let Value::Array(pair) = entry + && let [key, digest] = pair.as_slice() + { + let digest: Option = match digest { + Value::Null => None, + Value::Bool(_) + | Value::Number(_) + | Value::String(_) + | Value::Array(_) + | Value::Object(_) => Some(python_text(digest)), + }; + observed.insert(fact_key(key), digest); + } + } + observed +} diff --git a/crates/sqlbuild-rules/src/custom_host_pool/_helpers/merging.rs b/crates/sqlbuild-rules/src/custom_host_pool/_helpers/merging.rs new file mode 100644 index 0000000000..9df96b1bca --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/_helpers/merging.rs @@ -0,0 +1,116 @@ +//! Host slice merging steps. + +use std::collections::{BTreeMap, BTreeSet, HashMap}; + +use crate::custom_host_pool::constants::MAX_TRACKED_READS; +use crate::custom_host_pool::models::{HostEvaluation, HostSlice}; +use crate::custom_host_pool::types::{FactKey, HostInvocation}; + +/// Untrack evaluations of untracked rules, keeping each invocation's first position. +pub(crate) fn demoted( + evaluations: Vec, + untracked: &BTreeSet, + mut untracked_reads: BTreeMap>, +) -> (Vec, BTreeMap>) { + let mut evaluated: Vec = Vec::new(); + let mut positions: HashMap = HashMap::new(); + for mut evaluation in evaluations { + if untracked.contains(&evaluation.code) + && let Some(reads) = evaluation.reads.take() + { + untracked_reads + .entry(evaluation.code.clone()) + .or_default() + .extend(reads); + } + let key: HostInvocation = (evaluation.code.clone(), evaluation.subject.clone()); + match positions.get(&key) { + Some(index) => evaluated[*index] = evaluation, + None => { + positions.insert(key, evaluated.len()); + evaluated.push(evaluation); + } + } + } + (evaluated, untracked_reads) +} + +/// Each observed digest, forgetting any two hosts observed differently. +pub(crate) fn merged_observations(slices: &[HostSlice]) -> BTreeMap> { + let mut observed: BTreeMap> = BTreeMap::new(); + for slice in slices { + for (key, digest) in &slice.observed { + match observed.get(key) { + Some(seen) if seen != digest => { + observed.insert(key.clone(), None); + } + Some(_) => {} + None => { + observed.insert(key.clone(), digest.clone()); + } + } + } + } + observed +} + +/// Every slice's evaluations in single-host order, and the rules reported or demoted untracked. +pub(crate) fn ordered_evaluations( + slices: &[HostSlice], + order: &[HostInvocation], +) -> (Vec, BTreeSet) { + let positions: HashMap<&HostInvocation, usize> = order + .iter() + .enumerate() + .map(|(index, invocation)| (invocation, index)) + .collect(); + let mut evaluations: Vec = Vec::new(); + let mut untracked: BTreeSet = BTreeSet::new(); + for slice in slices { + evaluations.extend(slice.evaluations.iter().cloned()); + untracked.extend(slice.untracked_reads.keys().cloned()); + } + if slices.len() > 1 { + evaluations.sort_by_key(|evaluation| { + positions + .get(&(evaluation.code.clone(), evaluation.subject.clone())) + .copied() + .unwrap_or(order.len()) + }); + } + let tracked: BTreeSet<&Vec> = evaluations + .iter() + .filter(|evaluation| !untracked.contains(&evaluation.code)) + .filter_map(|evaluation| evaluation.reads.as_ref()) + .collect(); + let tracked_reads: usize = tracked.iter().map(|reads| reads.len()).sum(); + if slices.iter().any(|slice| slice.bounded) || tracked_reads > MAX_TRACKED_READS { + untracked.extend( + evaluations + .iter() + .filter(|evaluation| evaluation.reads.is_some()) + .map(|evaluation| evaluation.code.clone()), + ); + } + (evaluations, untracked) +} + +/// The reads each untracked rule's hosts reported directly. +pub(crate) fn collected_untracked_reads( + slices: &[HostSlice], + untracked: &BTreeSet, +) -> BTreeMap> { + let mut reads: BTreeMap> = untracked + .iter() + .map(|code| (code.clone(), BTreeSet::new())) + .collect(); + for slice in slices { + for (code, keys) in &slice.untracked_reads { + reads + .entry(code.clone()) + .or_default() + .extend(keys.iter().cloned()); + } + } + reads +} diff --git a/crates/sqlbuild-rules/src/custom_host_pool/_helpers/mod.rs b/crates/sqlbuild-rules/src/custom_host_pool/_helpers/mod.rs new file mode 100644 index 0000000000..7282700dcb --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/_helpers/mod.rs @@ -0,0 +1,4 @@ +pub(crate) mod cgroups; +pub(crate) mod decoding; +pub(crate) mod merging; +pub(crate) mod payload_values; diff --git a/crates/sqlbuild-rules/src/custom_host_pool/_helpers/payload_values.rs b/crates/sqlbuild-rules/src/custom_host_pool/_helpers/payload_values.rs new file mode 100644 index 0000000000..261efa60e5 --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/_helpers/payload_values.rs @@ -0,0 +1,36 @@ +//! Python `str()` of host payload values. + +use serde_json::Value; + +use crate::custom_host_pool::types::FactKey; + +/// `str(value)` for the scalar JSON values a host writes. +pub(crate) fn python_text(value: &Value) -> String { + match value { + Value::String(text) => text.clone(), + Value::Null => "None".to_owned(), + Value::Bool(true) => "True".to_owned(), + Value::Bool(false) => "False".to_owned(), + Value::Number(_) | Value::Array(_) | Value::Object(_) => value.to_string(), + } +} + +/// `tuple(map(str, key))` for one recorded fact key. +pub(crate) fn fact_key(value: &Value) -> FactKey { + match value { + Value::Array(parts) => parts.iter().map(python_text).collect(), + Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) | Value::Object(_) => { + Vec::new() + } + } +} + +/// `tuple(tuple(map(str, key)) for key in value)` for one read set. +pub(crate) fn readset(value: &Value) -> Vec { + match value { + Value::Array(keys) => keys.iter().map(fact_key).collect(), + Value::Null | Value::Bool(_) | Value::Number(_) | Value::String(_) | Value::Object(_) => { + Vec::new() + } + } +} diff --git a/crates/sqlbuild-rules/src/custom_host_pool/constants.rs b/crates/sqlbuild-rules/src/custom_host_pool/constants.rs new file mode 100644 index 0000000000..aedde400c8 --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/constants.rs @@ -0,0 +1,27 @@ +//! Host fan-out limits, the project subject and cgroup locations. + +/// The subject of an invocation that evaluates a project rule. +pub const CUSTOM_RULE_PROJECT_SUBJECT: &str = ""; +/// How long one host may take. +pub const HOST_TIMEOUT_MILLIS: u64 = 120_000; +/// Invocations that justify starting one more host. +pub const MIN_INVOCATIONS_PER_HOST: usize = 64; +/// The most hosts one run starts. +pub const MAX_HOSTS: usize = 8; +/// Tracked read keys above which every tracked rule is treated as untracked. +pub const MAX_TRACKED_READS: usize = 100_000; +/// The mounted cgroup hierarchy. +pub const CGROUP_ROOT: &str = "/sys/fs/cgroup"; +/// This process's cgroup memberships. +pub const PROC_CGROUP: &str = "/proc/self/cgroup"; +/// The cgroup v1 CPU controller name. +pub(crate) const CGROUP_V1_CPU_CONTROLLER: &str = "cpu"; +/// Directory names a cgroup v1 CPU controller may be mounted under. +pub(crate) const CGROUP_V1_CPU_DIRECTORIES: [&str; 3] = ["cpu", "cpu,cpuacct", "cpuacct,cpu"]; + +/// The `RulesError` message of a malformed host payload. +pub(crate) const INVALID_PAYLOAD: &str = "custom rule host returned an invalid payload"; +/// The `RulesError` message of a malformed evaluation. +pub(crate) const INVALID_EVALUATION: &str = "custom rule host returned an invalid evaluation"; +/// The `RulesError` message of findings that are not a list. +pub(crate) const INVALID_FINDINGS: &str = "custom rule host returned invalid findings"; diff --git a/crates/sqlbuild-rules/src/custom_host_pool/main/available_cores.rs b/crates/sqlbuild-rules/src/custom_host_pool/main/available_cores.rs new file mode 100644 index 0000000000..c2b8a21b03 --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/main/available_cores.rs @@ -0,0 +1,17 @@ +//! CPUs available to this process, including container CPU quotas. + +use std::path::Path; + +use crate::custom_host_pool::_helpers::cgroups::cgroup_quotas; + +/// The process CPU count capped by any cgroup quota listed in `proc_cgroup`, at least one. +#[must_use] +pub fn available_cores(process_cpus: usize, (cgroup_root, proc_cgroup): (&Path, &Path)) -> usize { + let entries: Vec = std::fs::read_to_string(proc_cgroup) + .map(|text| text.lines().map(str::to_owned).collect()) + .unwrap_or_default(); + cgroup_quotas(cgroup_root, &entries) + .into_iter() + .fold(process_cpus.max(1), usize::min) + .max(1) +} diff --git a/crates/sqlbuild-rules/src/custom_host_pool/main/decode_host_response.rs b/crates/sqlbuild-rules/src/custom_host_pool/main/decode_host_response.rs new file mode 100644 index 0000000000..912307faed --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/main/decode_host_response.rs @@ -0,0 +1,11 @@ +//! Decode one host's JSON response. + +use serde_json::Value; + +use crate::custom_host_pool::_helpers::decoding::decoded_slice; +use crate::custom_host_pool::models::HostSlice; + +/// The slice a host response describes, or the `RulesError` message of a malformed one. +pub fn decode_host_response(response: &Value) -> Result { + decoded_slice(response) +} diff --git a/crates/sqlbuild-rules/src/custom_host_pool/main/host_count.rs b/crates/sqlbuild-rules/src/custom_host_pool/main/host_count.rs new file mode 100644 index 0000000000..f0eeea510a --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/main/host_count.rs @@ -0,0 +1,12 @@ +//! How many hosts a run is worth. + +use crate::custom_host_pool::constants::{MAX_HOSTS, MIN_INVOCATIONS_PER_HOST}; + +/// One host per 64 invocations, bounded by usable CPUs and eight, and at least one. +#[must_use] +pub fn host_count(invocations: usize, cores: usize) -> usize { + MAX_HOSTS + .min(cores) + .min(invocations / MIN_INVOCATIONS_PER_HOST) + .max(1) +} diff --git a/crates/sqlbuild-rules/src/custom_host_pool/main/merge_host_slices.rs b/crates/sqlbuild-rules/src/custom_host_pool/main/merge_host_slices.rs new file mode 100644 index 0000000000..282652d78d --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/main/merge_host_slices.rs @@ -0,0 +1,32 @@ +//! Combine host slices exactly as one host would have reported every invocation. + +use std::collections::{BTreeMap, BTreeSet}; + +use crate::custom_host_pool::_helpers::merging::{ + collected_untracked_reads, demoted, merged_observations, ordered_evaluations, +}; +use crate::custom_host_pool::models::{HostEvaluation, HostRun, HostSlice}; +use crate::custom_host_pool::types::{FactKey, HostInvocation}; + +/// Merge slices in single-host invocation `order`, demoting tracking that grew too large. +#[must_use] +pub fn merge_host_slices(slices: Vec, order: &[HostInvocation]) -> HostRun { + let (evaluations, untracked): (Vec, BTreeSet) = + ordered_evaluations(&slices, order); + let untracked_reads: BTreeMap> = + collected_untracked_reads(&slices, &untracked); + let (evaluated, untracked_reads): (Vec, BTreeMap>) = + demoted(evaluations, &untracked, untracked_reads); + HostRun { + evaluated, + untracked_reads: untracked_reads + .into_iter() + .map(|(code, reads)| (code, reads.into_iter().collect())) + .collect(), + uncacheable: slices + .iter() + .flat_map(|slice| slice.uncacheable.iter().cloned()) + .collect(), + observed: merged_observations(&slices), + } +} diff --git a/crates/sqlbuild-rules/src/custom_host_pool/main/mod.rs b/crates/sqlbuild-rules/src/custom_host_pool/main/mod.rs new file mode 100644 index 0000000000..14dbd95162 --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/main/mod.rs @@ -0,0 +1,7 @@ +pub mod available_cores; +pub mod decode_host_response; +pub mod host_count; +pub mod merge_host_slices; +pub mod partition_host_plan; +pub mod planned_invocations; +pub mod without_codes; diff --git a/crates/sqlbuild-rules/src/custom_host_pool/main/partition_host_plan.rs b/crates/sqlbuild-rules/src/custom_host_pool/main/partition_host_plan.rs new file mode 100644 index 0000000000..276e7cc27c --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/main/partition_host_plan.rs @@ -0,0 +1,47 @@ +//! Deal planned invocations across hosts. + +use crate::custom_host_pool::constants::CUSTOM_RULE_PROJECT_SUBJECT; +use crate::custom_host_pool::types::{HostInvocation, HostPlan}; + +/// Project invocations and whole model subjects dealt round-robin across at most `hosts` plans. +#[must_use] +pub fn partition_host_plan( + plan: &HostPlan, + invocations: &[HostInvocation], + hosts: usize, +) -> Vec { + let mut units: Vec> = Vec::new(); + for invocation in invocations { + let starts_unit: bool = invocation.1 == CUSTOM_RULE_PROJECT_SUBJECT + || units + .last() + .and_then(|unit| unit.first()) + .is_none_or(|first| first.1 != invocation.1); + if starts_unit { + units.push(Vec::new()); + } + if let Some(unit) = units.last_mut() { + unit.push(invocation); + } + } + let count: usize = hosts.min(units.len()).max(1); + if count == 1 { + return vec![plan.clone()]; + } + let mut partitions: Vec = (0..count).map(|_| empty_partition(plan)).collect(); + for (index, unit) in units.into_iter().enumerate() { + for (code, subject) in unit { + if let Some(Some(subjects)) = partitions[index % count].get_mut(code) { + subjects.push(subject.clone()); + } + } + } + partitions +} + +/// Every planned rule with no subjects yet. +fn empty_partition(plan: &HostPlan) -> HostPlan { + plan.keys() + .map(|code| (code.clone(), Some(Vec::new()))) + .collect() +} diff --git a/crates/sqlbuild-rules/src/custom_host_pool/main/planned_invocations.rs b/crates/sqlbuild-rules/src/custom_host_pool/main/planned_invocations.rs new file mode 100644 index 0000000000..6b79a189db --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/main/planned_invocations.rs @@ -0,0 +1,40 @@ +//! Planned `(rule, subject)` invocations in the order one host evaluates them. + +use crate::custom_host_pool::constants::CUSTOM_RULE_PROJECT_SUBJECT; +use crate::custom_host_pool::models::HostRuleSubject; +use crate::custom_host_pool::types::{HostInvocation, HostPlan}; + +/// Project rules first, then each model path's selected model rules, all in rule code order. +#[must_use] +pub fn planned_invocations( + plan: &HostPlan, + rules: &[(String, HostRuleSubject)], + model_paths: &[String], +) -> Vec { + let mut planned: Vec<&(String, HostRuleSubject)> = rules + .iter() + .filter(|(code, _)| plan.contains_key(code)) + .collect(); + planned.sort_by(|left, right| left.0.cmp(&right.0)); + let mut invocations: Vec = planned + .iter() + .filter(|(_, subject)| *subject == HostRuleSubject::Project) + .map(|(code, _)| (code.clone(), CUSTOM_RULE_PROJECT_SUBJECT.to_owned())) + .collect(); + let model_rules: Vec<(&String, Option<&Vec>)> = planned + .iter() + .filter(|(_, subject)| *subject == HostRuleSubject::Model) + .map(|(code, _)| (code, plan.get(code).and_then(Option::as_ref))) + .collect(); + if model_rules.is_empty() { + return invocations; + } + for path in model_paths { + for (code, selected) in &model_rules { + if selected.is_none_or(|subjects| subjects.contains(path)) { + invocations.push(((*code).clone(), path.clone())); + } + } + } + invocations +} diff --git a/crates/sqlbuild-rules/src/custom_host_pool/main/without_codes.rs b/crates/sqlbuild-rules/src/custom_host_pool/main/without_codes.rs new file mode 100644 index 0000000000..c75fb8da37 --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/main/without_codes.rs @@ -0,0 +1,25 @@ +//! Drop every result a host reported for some rules. + +use std::collections::BTreeSet; + +use crate::custom_host_pool::models::HostSlice; + +/// The slice without evaluations, untracked reads, uncacheable or stateful marks of `codes`. +#[must_use] +pub fn without_codes(slice: HostSlice, codes: &BTreeSet) -> HostSlice { + HostSlice { + evaluations: slice + .evaluations + .into_iter() + .filter(|evaluation| !codes.contains(&evaluation.code)) + .collect(), + untracked_reads: slice + .untracked_reads + .into_iter() + .filter(|(code, _)| !codes.contains(code)) + .collect(), + uncacheable: slice.uncacheable.difference(codes).cloned().collect(), + stateful: slice.stateful.difference(codes).cloned().collect(), + ..slice + } +} diff --git a/crates/sqlbuild-rules/src/custom_host_pool/mod.rs b/crates/sqlbuild-rules/src/custom_host_pool/mod.rs new file mode 100644 index 0000000000..753b8bddaf --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/mod.rs @@ -0,0 +1,10 @@ +//! Custom-rule host fan-out: invocation order, host partitions, CPU capacity and merged results. + +pub(crate) mod _helpers; +pub mod constants; +pub mod main; +pub mod models; +pub mod types; + +#[cfg(test)] +mod tests; diff --git a/crates/sqlbuild-rules/src/custom_host_pool/models.rs b/crates/sqlbuild-rules/src/custom_host_pool/models.rs new file mode 100644 index 0000000000..a9cd86f651 --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/models.rs @@ -0,0 +1,42 @@ +//! One host's reported results and the results merged across hosts. + +use std::collections::{BTreeMap, BTreeSet}; + +use crate::custom_host_pool::types::FactKey; + +/// The kind of subject a custom rule evaluates. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum HostRuleSubject { + Project, + Model, + Other, +} + +/// One evaluated invocation: its findings and, when tracked, the facts it read. +#[derive(Clone, Debug, PartialEq)] +pub struct HostEvaluation { + pub code: String, + pub subject: String, + pub findings: Vec, + pub reads: Option>, +} + +/// What one host reported for its share of the plan. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct HostSlice { + pub evaluations: Vec, + pub untracked_reads: BTreeMap>, + pub uncacheable: BTreeSet, + pub observed: BTreeMap>, + pub bounded: bool, + pub stateful: BTreeSet, +} + +/// Every host's results combined exactly as one host would have reported them. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct HostRun { + pub evaluated: Vec, + pub untracked_reads: BTreeMap>, + pub uncacheable: BTreeSet, + pub observed: BTreeMap>, +} diff --git a/crates/sqlbuild-rules/src/custom_host_pool/tests.rs b/crates/sqlbuild-rules/src/custom_host_pool/tests.rs new file mode 100644 index 0000000000..f2aafd9105 --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/tests.rs @@ -0,0 +1,6 @@ +#[path = "tests/test_custom_host_pool.rs"] +mod custom_host_pool; +#[path = "tests/helpers.rs"] +mod helpers; +#[path = "tests/test_types.rs"] +mod test_types; diff --git a/crates/sqlbuild-rules/src/custom_host_pool/tests/helpers.rs b/crates/sqlbuild-rules/src/custom_host_pool/tests/helpers.rs new file mode 100644 index 0000000000..1340930d24 --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/tests/helpers.rs @@ -0,0 +1,56 @@ +use std::path::{Path, PathBuf}; + +use crate::custom_host_pool::models::HostEvaluation; +use crate::custom_host_pool::types::{HostInvocation, HostPlan}; + +/// Write a fake `/proc/self/cgroup` and a cgroup hierarchy below `root`. +pub(super) fn write_cgroup_files( + root: &Path, + (proc_cgroup, files): (&[u8], &[(&str, &[u8])]), +) -> PathBuf { + for (relative_path, contents) in files { + let path: PathBuf = root.join("cgroup").join(relative_path); + std::fs::create_dir_all(path.parent().expect("parent")).expect("cgroup folder"); + std::fs::write(path, contents).expect("cgroup file"); + } + let proc: PathBuf = root.join("proc-self-cgroup"); + std::fs::write(&proc, proc_cgroup).expect("process cgroup"); + proc +} + +pub(super) fn plan(entries: &[(&str, Option<&[&str]>)]) -> HostPlan { + entries + .iter() + .map(|(code, subjects)| { + ( + (*code).to_owned(), + subjects.map(|subjects| { + subjects + .iter() + .map(|subject| (*subject).to_owned()) + .collect() + }), + ) + }) + .collect() +} + +pub(super) fn invocations(entries: &[(&str, &str)]) -> Vec { + entries + .iter() + .map(|(code, subject)| ((*code).to_owned(), (*subject).to_owned())) + .collect() +} + +pub(super) fn evaluation(code: &str, subject: &str, reads: Option<&[&[&str]]>) -> HostEvaluation { + HostEvaluation { + code: code.to_owned(), + subject: subject.to_owned(), + findings: Vec::new(), + reads: reads.map(|keys| { + keys.iter() + .map(|key| key.iter().map(|part| (*part).to_owned()).collect()) + .collect() + }), + } +} diff --git a/crates/sqlbuild-rules/src/custom_host_pool/tests/test_custom_host_pool.rs b/crates/sqlbuild-rules/src/custom_host_pool/tests/test_custom_host_pool.rs new file mode 100644 index 0000000000..84a6c451e1 --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/tests/test_custom_host_pool.rs @@ -0,0 +1,294 @@ +use std::collections::{BTreeMap, BTreeSet}; +use std::path::PathBuf; + +use serde_json::json; + +use crate::custom_host_pool::main::available_cores::available_cores; +use crate::custom_host_pool::main::decode_host_response::decode_host_response; +use crate::custom_host_pool::main::host_count::host_count; +use crate::custom_host_pool::main::merge_host_slices::merge_host_slices; +use crate::custom_host_pool::main::partition_host_plan::partition_host_plan; +use crate::custom_host_pool::models::HostSlice; +use crate::custom_host_pool::tests::helpers::{evaluation, invocations, plan, write_cgroup_files}; +use crate::custom_host_pool::tests::test_types::{ + DecodeTestCase, HostCapacityTestCase, HostCountTestCase, HostPlanTestCase, MergeTestCase, +}; + +const NODE_CPUS: usize = 64; + +#[test] +fn given_cgroup_quota_when_counting_cores_then_quota_caps_process_cpus() { + let test_cases = [ + HostCapacityTestCase { + description: "cgroup v2 quota on the process cgroup", + proc_cgroup: b"0::/kubepods/orders\n", + files: &[("kubepods/orders/cpu.max", b"400000 100000\n")], + expected_cores: 4, + }, + HostCapacityTestCase { + description: "cgroup v2 fractional quota rounds up", + proc_cgroup: b"0::/kubepods/orders\n", + files: &[("kubepods/orders/cpu.max", b"150000 100000\n")], + expected_cores: 2, + }, + HostCapacityTestCase { + description: "cgroup v2 ancestor quota is tighter", + proc_cgroup: b"0::/kubepods/orders\n", + files: &[ + ("kubepods/cpu.max", b"200000 100000\n"), + ("kubepods/orders/cpu.max", b"800000 100000\n"), + ], + expected_cores: 2, + }, + HostCapacityTestCase { + description: "cgroup v2 without a quota", + proc_cgroup: b"0::/\n", + files: &[("cpu.max", b"max 100000\n")], + expected_cores: NODE_CPUS, + }, + HostCapacityTestCase { + description: "cgroup v1 CFS quota", + proc_cgroup: b"4:cpu,cpuacct:/docker/orders\n", + files: &[ + ("cpu,cpuacct/docker/orders/cpu.cfs_quota_us", b"300000\n"), + ("cpu,cpuacct/docker/orders/cpu.cfs_period_us", b"100000\n"), + ], + expected_cores: 3, + }, + HostCapacityTestCase { + description: "cgroup v1 unlimited CFS quota", + proc_cgroup: b"4:cpu,cpuacct:/\n", + files: &[ + ("cpu,cpuacct/cpu.cfs_quota_us", b"-1\n"), + ("cpu,cpuacct/cpu.cfs_period_us", b"100000\n"), + ], + expected_cores: NODE_CPUS, + }, + HostCapacityTestCase { + description: "no cgroup files", + proc_cgroup: b"", + files: &[], + expected_cores: NODE_CPUS, + }, + HostCapacityTestCase { + description: "non-UTF-8 cgroup v2 quota is ignored", + proc_cgroup: b"0::/kubepods/orders\n", + files: &[("kubepods/orders/cpu.max", b"\xff\xfe 100000\n")], + expected_cores: NODE_CPUS, + }, + HostCapacityTestCase { + description: "non-UTF-8 process cgroup is ignored", + proc_cgroup: b"0::/kubepods/\xff\n", + files: &[("kubepods/cpu.max", b"200000 100000\n")], + expected_cores: NODE_CPUS, + }, + ]; + for test_case in test_cases { + let directory = tempfile::tempdir().expect("a temporary folder"); + let proc: PathBuf = + write_cgroup_files(directory.path(), (test_case.proc_cgroup, test_case.files)); + + let cores: usize = available_cores(NODE_CPUS, (&directory.path().join("cgroup"), &proc)); + + assert_eq!(cores, test_case.expected_cores, "{}", test_case.description); + } +} + +#[test] +fn given_planned_invocations_when_partitioning_then_split_is_deterministic() { + let test_cases = [ + HostPlanTestCase { + description: "project rules and whole subjects are dealt round-robin", + plan: plan(&[ + ("XSQBRT101", None), + ("XSQBRT102", None), + ("XSQBRT103", Some(&["models/b.sql"])), + ]), + invocations: invocations(&[ + ("XSQBRT102", ""), + ("XSQBRT101", "models/a.sql"), + ("XSQBRT101", "models/b.sql"), + ("XSQBRT103", "models/b.sql"), + ("XSQBRT101", "models/c.sql"), + ]), + hosts: 2, + expected_plans: vec![ + plan(&[ + ("XSQBRT101", Some(&["models/b.sql"])), + ("XSQBRT102", Some(&[""])), + ("XSQBRT103", Some(&["models/b.sql"])), + ]), + plan(&[ + ("XSQBRT101", Some(&["models/a.sql", "models/c.sql"])), + ("XSQBRT102", Some(&[])), + ("XSQBRT103", Some(&[])), + ]), + ], + }, + HostPlanTestCase { + description: "one subject keeps the original plan", + plan: plan(&[("XSQBRT101", None)]), + invocations: invocations(&[("XSQBRT101", "models/a.sql")]), + hosts: 4, + expected_plans: vec![plan(&[("XSQBRT101", None)])], + }, + ]; + for test_case in test_cases { + let plans = partition_host_plan(&test_case.plan, &test_case.invocations, test_case.hosts); + + assert_eq!(plans, test_case.expected_plans, "{}", test_case.description); + } +} + +#[test] +fn given_invocation_count_when_sizing_hosts_then_hosts_follow_work_and_cores() { + let test_cases = [ + HostCountTestCase { + description: "hosts are capped on a large node", + invocations: 10_000, + cores: NODE_CPUS, + expected_hosts: 8, + }, + HostCountTestCase { + description: "few invocations run on one host", + invocations: 63, + cores: NODE_CPUS, + expected_hosts: 1, + }, + HostCountTestCase { + description: "cores cap the hosts", + invocations: 1_000, + cores: 2, + expected_hosts: 2, + }, + ]; + for test_case in test_cases { + assert_eq!( + host_count(test_case.invocations, test_case.cores), + test_case.expected_hosts, + "{}", + test_case.description + ); + } +} + +#[test] +fn given_host_response_when_decoding_then_slice_or_rules_error_is_returned() { + let test_cases = [ + DecodeTestCase { + description: "tracked reads come from the read set table", + response: json!({ + "readsets": [[["project.models"]]], + "evaluations": [ + {"code": "R1", "subject": null, "findings": [{"message": "m"}], "reads": 0}, + {"code": "R2", "subject": "models/a.sql", "findings": [], "reads": 0}, + ], + "untracked": {"R2": [["tree.paths"]]}, + "uncacheable": ["R3"], + "observed": [[["tree.paths"], "abc"], [["tree.glob", "*.sql"], null]], + "bounded": true, + "stateful": ["R1"], + }), + expected_slice: Ok(HostSlice { + evaluations: vec![ + crate::custom_host_pool::models::HostEvaluation { + findings: vec![json!({"message": "m"})], + ..evaluation("R1", "", Some(&[&["project.models"]])) + }, + evaluation("R2", "models/a.sql", None), + ], + untracked_reads: BTreeMap::from([( + "R2".to_owned(), + vec![vec!["tree.paths".to_owned()]], + )]), + uncacheable: BTreeSet::from(["R3".to_owned()]), + observed: BTreeMap::from([ + (vec!["tree.glob".to_owned(), "*.sql".to_owned()], None), + (vec!["tree.paths".to_owned()], Some("abc".to_owned())), + ]), + bounded: true, + stateful: BTreeSet::from(["R1".to_owned()]), + }), + }, + DecodeTestCase { + description: "a payload without its tables is invalid", + response: json!({"readsets": []}), + expected_slice: Err("custom rule host returned an invalid payload".to_owned()), + }, + DecodeTestCase { + description: "findings must be a list", + response: json!({ + "readsets": [], "untracked": {}, "uncacheable": [], "observed": [], + "evaluations": [{"code": "R1", "findings": {}}], + }), + expected_slice: Err("custom rule host returned invalid findings".to_owned()), + }, + ]; + for test_case in test_cases { + assert_eq!( + decode_host_response(&test_case.response), + test_case.expected_slice, + "{}", + test_case.description + ); + } +} + +#[test] +fn given_host_slices_when_merging_then_results_follow_single_host_order() { + let test_cases = [ + MergeTestCase { + description: "split slices are reordered and conflicting observations forgotten", + slices: vec![ + HostSlice { + evaluations: vec![evaluation("R1", "b", Some(&[&["k"]]))], + observed: BTreeMap::from([(vec!["t".to_owned()], Some("1".to_owned()))]), + ..HostSlice::default() + }, + HostSlice { + evaluations: vec![evaluation("R1", "a", Some(&[&["k"]]))], + observed: BTreeMap::from([(vec!["t".to_owned()], Some("2".to_owned()))]), + ..HostSlice::default() + }, + ], + order: invocations(&[("R1", "a"), ("R1", "b")]), + expected_order: invocations(&[("R1", "a"), ("R1", "b")]), + expected_untracked: vec![], + expected_observed: vec![(vec!["t".to_owned()], None)], + }, + MergeTestCase { + description: "a bounded host untracks every tracked rule", + slices: vec![HostSlice { + evaluations: vec![evaluation("R1", "a", Some(&[&["k"]]))], + bounded: true, + ..HostSlice::default() + }], + order: invocations(&[("R1", "a")]), + expected_order: invocations(&[("R1", "a")]), + expected_untracked: vec!["R1".to_owned()], + expected_observed: vec![], + }, + ]; + for test_case in test_cases { + let run = merge_host_slices(test_case.slices, &test_case.order); + let order: Vec<(String, String)> = run + .evaluated + .iter() + .map(|evaluation| (evaluation.code.clone(), evaluation.subject.clone())) + .collect(); + + assert_eq!(order, test_case.expected_order, "{}", test_case.description); + assert_eq!( + run.untracked_reads.keys().cloned().collect::>(), + test_case.expected_untracked, + "{}", + test_case.description + ); + assert_eq!( + run.observed.into_iter().collect::>(), + test_case.expected_observed, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-rules/src/custom_host_pool/tests/test_types.rs b/crates/sqlbuild-rules/src/custom_host_pool/tests/test_types.rs new file mode 100644 index 0000000000..032f6e8e11 --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/tests/test_types.rs @@ -0,0 +1,39 @@ +use crate::custom_host_pool::models::HostSlice; +use crate::custom_host_pool::types::{HostInvocation, HostPlan}; + +pub(super) struct HostCapacityTestCase { + pub(super) description: &'static str, + pub(super) proc_cgroup: &'static [u8], + pub(super) files: &'static [(&'static str, &'static [u8])], + pub(super) expected_cores: usize, +} + +pub(super) struct HostPlanTestCase { + pub(super) description: &'static str, + pub(super) plan: HostPlan, + pub(super) invocations: Vec, + pub(super) hosts: usize, + pub(super) expected_plans: Vec, +} + +pub(super) struct HostCountTestCase { + pub(super) description: &'static str, + pub(super) invocations: usize, + pub(super) cores: usize, + pub(super) expected_hosts: usize, +} + +pub(super) struct DecodeTestCase { + pub(super) description: &'static str, + pub(super) response: serde_json::Value, + pub(super) expected_slice: Result, +} + +pub(super) struct MergeTestCase { + pub(super) description: &'static str, + pub(super) slices: Vec, + pub(super) order: Vec, + pub(super) expected_order: Vec, + pub(super) expected_untracked: Vec, + pub(super) expected_observed: Vec<(Vec, Option)>, +} diff --git a/crates/sqlbuild-rules/src/custom_host_pool/types.rs b/crates/sqlbuild-rules/src/custom_host_pool/types.rs new file mode 100644 index 0000000000..539859fd43 --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/types.rs @@ -0,0 +1,10 @@ +//! Host plans and invocations. + +use std::collections::BTreeMap; + +/// Subjects to evaluate per rule code; `None` evaluates every subject. +pub type HostPlan = BTreeMap>>; +/// One `(rule code, subject)` evaluation. +pub type HostInvocation = (String, String); +/// One recorded fact read. +pub type FactKey = Vec; diff --git a/crates/sqlbuild-rules/src/lib.rs b/crates/sqlbuild-rules/src/lib.rs index d6eb3c2b9b..7c4bd3a4cb 100644 --- a/crates/sqlbuild-rules/src/lib.rs +++ b/crates/sqlbuild-rules/src/lib.rs @@ -4,6 +4,7 @@ pub mod configuration; pub mod constants; +pub mod custom_host_pool; pub mod custom_rule_cache; pub mod dsl_headers; pub mod engine; diff --git a/python-boundary.toml b/python-boundary.toml index c8f7239217..6e3e23603d 100644 --- a/python-boundary.toml +++ b/python-boundary.toml @@ -2,7 +2,7 @@ # migration work that moves to Rust. `remaining_lines` must equal the current count of that work: # a PR that moves code lowers it, and adding movable Python fails the check. # Native failures must surface as errors; `handled_native_errors` lists the reviewed exceptions. -remaining_lines = 9599 +remaining_lines = 9564 [[stays]] path = "executor" @@ -112,12 +112,7 @@ note = "Samples process resources for debug logging through Python's resource mo [[stays]] path = "rule_engine/_helpers/host" reason = "user-code" -note = "Subprocess host that loads and runs user rule modules." - -[[stays]] -path = "rule_engine/_helpers/engine/custom_dispatch.py" -reason = "user-code" -note = "Submits custom rules to their worker thread beside native built-in rule evaluation." +note = "Rule 2: the subprocess host that loads and runs user rule modules (custom_host, custom_evaluation, tracked_facts, module_state, custom_host_launch); host fan-out, partitioning and CPU capacity run natively." [[stays]] path = "rule_engine/_helpers/engine/custom_payloads.py" @@ -129,16 +124,6 @@ path = "rule_engine/_helpers/engine/_custom_catalogue.py" reason = "user-code" note = "Imports the project's custom rule modules and collects their @rule functions for the native catalogue." -[[stays]] -path = "rule_engine/_helpers/engine/custom_rules.py" -reason = "user-code" -note = "Dispatches custom rules to the rule-host worker pool and digests the fact views they read." - -[[stays]] -path = "rule_engine/_helpers/run/rule_threads.py" -reason = "user-code" -note = "Runs rule phases and early SQL lint on worker threads with futures and stop events." - [[stays]] path = "rule_engine/classes/rule_context.py" reason = "user-code" @@ -162,7 +147,7 @@ note = "Public custom-rule test harness that runs user rules over a real compile [[stays]] path = "rule_engine/_helpers/engine/fact_replay.py" reason = "user-code" -note = "Recomputes the fact-view values a custom rule read, for its cache key." +note = "Rule 2: the fact digest scheme shared with the rule host (tracked_facts.py, custom_host.py); recomputes values through the user-facing fact views and encodes live Python objects with orjson." [[stays]] path = "rule_engine/main/define.py" @@ -174,11 +159,6 @@ path = "rule_engine/main/evaluate_rule.py" reason = "user-code" note = "Public custom-rule test harness entry point." -[[stays]] -path = "rule_engine/main/_evaluate.py" -reason = "user-code" -note = "Project evaluation entry used by the public custom-rule test harness." - [[stays]] path = "rule_engine/classes/audit_facts.py" reason = "user-code" @@ -234,11 +214,6 @@ path = "rule_engine/classes/fact_reads.py" reason = "user-code" note = "Records the fact-view reads of one user rule evaluation." -[[stays]] -path = "rule_engine/classes/fact_digests.py" -reason = "user-code" -note = "Digests the fact-view values user rules read, for the custom-rule cache." - [[stays]] path = "rule_engine/classes/shared_fact_encoder.py" reason = "user-code" diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index 2902ab48ef..0ad4c38c75 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -2,7 +2,7 @@ import threading from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence, Set -from concurrent.futures import Future +from concurrent.futures import Executor, Future from contextlib import AbstractContextManager from contextvars import Token from dataclasses import Field @@ -351,9 +351,11 @@ from sqlbuild.presentation.types import TransientLineOwner from sqlbuild.provider.classes.container import ProviderContainer from sqlbuild.python_nodes.models import SqlResourceRef from sqlbuild.rule_engine.models import ( + CustomRulesOutcome, Finding, NativeRulesEvaluation, PreparedSqlLint, + PreparedSqlLintResult, Rule, RuleRunPlan, RulesConfig, @@ -6500,38 +6502,59 @@ def custom_rule_identity( dialect: str, project_dir: str, ) -> str: ... - -class CustomRuleCache: - """One custom-rule run's subjects, reused and new results, and its cache file.""" - - def __init__( - self, - *, - path: Path | None, - rules: list[tuple[str, str | None, str | None]], - target_name: str | None, - models: list[tuple[str, str, str | None]], - ) -> None: ... - def pending_readsets(self) -> list[tuple[tuple[str, ...], ...]]: ... - def plan(self, *, digests: list[str | None]) -> dict[str, list[str] | None]: ... - def record( - self, - *, - evaluated: list[ - tuple[ - str, - str, - tuple[tuple[tuple[str, ...], ...], str | None] | None, - list[dict[str, object]], - ] - ], - ) -> None: ... - def findings(self) -> list[Finding]: ... - @property - def hits(self) -> int: ... - @property - def misses(self) -> int: ... - +def evaluate_project_rules( + *, + project: CompiledProjectView, + config: RulesConfig, + project_dir: Path, + dialect: str = "generic", + verify_determinism: bool = False, +) -> RulesResult: ... +def evaluate_native_rules( + *, + project: CompiledProjectView, + config: RulesConfig, + project_dir: Path, + catalogue: tuple[Rule, ...], + dialect: str = "generic", + initial_findings: tuple[Finding, ...] | None = None, + defer_suppressions: bool = False, + verify_determinism: bool = False, +) -> RulesResult: ... +def evaluate_custom_rules( + *, + project: CompiledProjectView, + config: RulesConfig, + project_dir: Path, + rules: tuple[Rule, ...], + dialect: str, + verify_determinism: bool = False, + implementation_fingerprints: dict[str, str] | None = None, +) -> CustomRulesOutcome: ... +def evaluate_rules( + *, + graph: ProjectGraph, + discovered_inputs: DiscoveredProjectInputs, + config: RulesConfig, + project_dir: Path, + dialect: str, + selected_keys: frozenset[CompiledObjectKey] | None = None, + prepared_sql: PreparedSqlLint | None = None, + expansion_reuse: SqlExpansionReuse | None = None, + no_sql_analysis: bool = False, +) -> RulesRunResult: ... +def prepare_sql_rules( + *, inputs: CompileInputsView, executor: Executor, dialect: str, stop: Event +) -> PreparedSqlLint | None: ... +def prepare_early_sql_lint( + *, + project_dir: Path, + config: LintConfig, + discovered_inputs: DiscoveredProjectInputs, + static_declaration_scope: DeclarationScopeBuild | None, + compiled_expansions: dict[Path, CompiledSqlExpansion], + stop: Event, +) -> PreparedSqlLintResult: ... def custom_rule_test_evidence( *, rule: Rule, project_dir: Path ) -> list[tuple[str, int, int, str]]: ... diff --git a/src/sqlbuild/rule_engine/_helpers/engine/custom_dispatch.py b/src/sqlbuild/rule_engine/_helpers/engine/custom_dispatch.py deleted file mode 100644 index 1f5b2a093e..0000000000 --- a/src/sqlbuild/rule_engine/_helpers/engine/custom_dispatch.py +++ /dev/null @@ -1,177 +0,0 @@ -"""Dispatch custom rules to their worker thread beside native built-in rule evaluation.""" - -from __future__ import annotations - -from collections.abc import Callable -from concurrent.futures import Executor, Future, ThreadPoolExecutor -from functools import partial -from pathlib import Path - -import sqlbuild._native as _native -from sqlbuild._native import ( - build_rules_request, - evaluate_rules_request, - finalize_rule_findings_rows, - rule_selected, - selected_rule_codes, -) -from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.rule_engine._helpers.engine.custom_payloads import custom_rule_payloads -from sqlbuild.rule_engine._helpers.engine.custom_rules import evaluate_custom_rules_cached -from sqlbuild.rule_engine.constants import TYPE_PROOF_RULE_CODES -from sqlbuild.rule_engine.models import ( - CustomRulesOutcome, - Finding, - NativeRulesEvaluation, - Rule, - RulesConfig, - RulesResult, -) - - -def evaluate_native( - *, - project: CompiledProjectView, - config: RulesConfig, - project_dir: Path, - catalogue: tuple[Rule, ...], - dialect: str = "generic", - initial_findings: tuple[Finding, ...] = (), - defer_suppressions: bool = False, - verify_determinism: bool = False, - custom_payloads: list[dict[str, object]] | None = None, - custom_outcome: Future[CustomRulesOutcome] | None = None, -) -> RulesResult: - """Evaluate built-in rules natively while custom rules run incrementally beside them.""" - - payloads: list[dict[str, object]] = ( - custom_rule_payloads(catalogue=catalogue, project_dir=project_dir) - if custom_payloads is None - else custom_payloads - ) - return _evaluate_native_rows( - project=project, - config=config, - project_dir=project_dir, - dialect=dialect, - initial_findings=initial_findings, - defer_suppressions=defer_suppressions, - custom_payloads=payloads, - start_custom=( - (lambda **_: custom_outcome) - if custom_payloads is not None - else partial( - start_custom_rules, - project=project, - config=config, - project_dir=project_dir, - catalogue=catalogue, - custom_payloads=payloads, - dialect=dialect, - verify_determinism=verify_determinism, - ) - ), - ) - - -def _evaluate_native_rows( - *, - project: CompiledProjectView, - config: RulesConfig, - project_dir: Path, - dialect: str, - initial_findings: tuple[Finding, ...], - defer_suppressions: bool, - custom_payloads: list[dict[str, object]], - start_custom: Callable[..., Future[CustomRulesOutcome] | None], -) -> RulesResult: - request: _native.NativeRulesRequest = build_rules_request( - project=project, - config=config, - project_dir=project_dir, - dialect=dialect, - initial_findings=initial_findings, - custom_payloads=custom_payloads, - include_type_proof=any( - rule_selected(config=config, code=code) for code in TYPE_PROOF_RULE_CODES - ), - ) - with ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-custom-rules") as pool: - custom_future: Future[CustomRulesOutcome] | None = start_custom(executor=pool) - native: NativeRulesEvaluation = evaluate_rules_request(request) - del request - custom: CustomRulesOutcome = ( - CustomRulesOutcome(findings=(), cache_hits=0, cache_misses=0, custom_ms=0) - if custom_future is None - else custom_future.result() - ) - findings: tuple[Finding, ...] = ( - *native.findings, - *custom.findings, - ) - if not defer_suppressions: - findings = finalize_rule_findings_rows( - config=config, - project_dir=project_dir, - evaluated_codes=native.selected_codes, - findings=findings, - ) - return RulesResult( - findings=findings, - evaluated_models=native.evaluated_models, - cache_hits=(native.cache_hits + native.cache_misses if native.reused else native.cache_hits) - + custom.cache_hits, - cache_misses=(0 if native.reused else native.cache_misses) + custom.cache_misses, - built_in_ms=0 if native.reused else native.built_in_ms, - custom_ms=custom.custom_ms, - custom_cpu_ms=custom.custom_cpu_ms, - ) - - -def start_custom_rules( - *, - executor: Executor, - project: CompiledProjectView, - config: RulesConfig, - project_dir: Path, - catalogue: tuple[Rule, ...], - custom_payloads: list[dict[str, object]], - dialect: str, - verify_determinism: bool = False, -) -> Future[CustomRulesOutcome] | None: - """Submit selected custom rules after digesting installed code here, clear of GIL contention.""" - - custom_rules: tuple[Rule, ...] = _selected_custom_rules( - config=config, catalogue=catalogue, custom_payloads=custom_payloads - ) - if not custom_rules: - return None - if config.cache.enabled: - _ = _native.compiled_code_identity() - return executor.submit( - evaluate_custom_rules_cached, - project=project, - config=config, - project_dir=project_dir.resolve(), - rules=custom_rules, - dialect=dialect, - verify_determinism=verify_determinism, - implementation_fingerprints={ - str(payload["code"]): str(payload["implementation_fingerprint"]) - for payload in custom_payloads - }, - ) - - -def _selected_custom_rules( - *, - config: RulesConfig, - catalogue: tuple[Rule, ...], - custom_payloads: list[dict[str, object]], -) -> tuple[Rule, ...]: - if not custom_payloads: - return () - selected: frozenset[str] = frozenset( - selected_rule_codes(config=config, custom_payloads=custom_payloads) - ) - return tuple(rule for rule in catalogue if rule.custom and rule.code in selected) diff --git a/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py b/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py deleted file mode 100644 index 4f6af8389d..0000000000 --- a/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py +++ /dev/null @@ -1,233 +0,0 @@ -"""Incremental custom-rule evaluation keyed on the compiler facts each invocation read.""" - -from __future__ import annotations - -import os -import time -from collections.abc import Callable -from dataclasses import replace -from pathlib import Path -from typing import Any - -import sqlbuild._native as _native -from sqlbuild._native import ( - CustomRuleCache, - compiled_project_record, - custom_rule_identity, - custom_rule_implementation_fingerprint, - custom_rule_import_closure, - rules_bulk_cache_path, -) -from sqlbuild.compiler.compile.models import CompiledProject -from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.compiler.scopes.models import ScopeIndex -from sqlbuild.rule_engine._helpers.engine.fact_replay import full_fact_keys, is_tree_fact -from sqlbuild.rule_engine._helpers.host.custom_host_pool import run_custom_hosts -from sqlbuild.rule_engine.classes.fact_digests import FactDigests -from sqlbuild.rule_engine.classes.rule_context import RuleFactViews, build_rule_fact_views -from sqlbuild.rule_engine.constants import ( - CUSTOM_HOST_RUNTIME_VERSION, - CUSTOM_RULES_CACHE_FILE, -) -from sqlbuild.rule_engine.exceptions import FactDigestError -from sqlbuild.rule_engine.models import CustomHostRun, CustomRulesOutcome, Rule, RulesConfig -from sqlbuild.rule_engine.types import CustomEvaluationRow, FactKey - - -def evaluate_custom_rules_cached( - *, - project: CompiledProjectView, - config: RulesConfig, - project_dir: Path, - rules: tuple[Rule, ...], - dialect: str, - verify_determinism: bool = False, - implementation_fingerprints: dict[str, str] | None = None, -) -> CustomRulesOutcome: - """Reuse results whose recorded fact reads are unchanged and evaluate only the rest.""" - - started: float = time.monotonic() - started_cpu: float = _cpu_seconds() - custom: tuple[Rule, ...] = tuple(sorted(rules, key=lambda rule: rule.code)) - if not custom: - return CustomRulesOutcome(findings=(), cache_hits=0, cache_misses=0, custom_ms=0) - facts_project: CompiledProjectView = host_project(project) - enabled: bool = config.cache.enabled - cache: CustomRuleCache = CustomRuleCache( - path=( - rules_bulk_cache_path(project_dir=project_dir, file_name=CUSTOM_RULES_CACHE_FILE) - if enabled - else None - ), - rules=[ - ( - rule.code, - _rule_identity( - rule=rule, - config=config, - project_dir=project_dir, - dialect=dialect, - fingerprint=(implementation_fingerprints or {}).get(rule.code), - ) - if enabled - else None, - None if rule.subject is None else rule.subject.value, - ) - for rule in custom - ], - target_name=project.effective_target_name, - models=[ - (model.relative_path.as_posix(), model.name, _materialization(model)) - for model in project.models - ], - ) - digests: FactDigests = FactDigests( - views=_views_factory(project=facts_project, project_dir=project_dir, dialect=dialect) - ) - plan: dict[str, list[str] | None] = cache.plan( - digests=[ - _readset_digest(readset=readset, digests=digests) - for readset in cache.pending_readsets() - ] - ) - if plan: - run: CustomHostRun = run_custom_hosts( - project=facts_project, - config=config, - project_dir=project_dir, - dialect=dialect, - rules=custom, - model_paths=tuple(sorted(model.relative_path.as_posix() for model in project.models)), - plan=plan, - track_reads=enabled, - verify_determinism=verify_determinism, - ) - cache.record(evaluated=_evaluations(run=run, digests=digests, cacheable=enabled)) - return CustomRulesOutcome( - findings=tuple(cache.findings()), - cache_hits=cache.hits, - cache_misses=cache.misses, - custom_ms=_elapsed_ms(started), - custom_cpu_ms=_cpu_ms(started_cpu), - ) - - -def _views_factory( - *, project: CompiledProjectView, project_dir: Path, dialect: str -) -> Callable[[], RuleFactViews]: - def build() -> RuleFactViews: - return build_rule_fact_views(project=project, project_dir=project_dir, dialect=dialect) - - return build - - -def _materialization(model: Any) -> str | None: - value: object = model.config.values.get("materialized") - return value if isinstance(value, str) else None - - -def _rule_identity( - *, rule: Rule, config: RulesConfig, project_dir: Path, dialect: str, fingerprint: str | None -) -> str: - return custom_rule_identity( - code_identity=_native.compiled_code_identity(), - runtime_version=CUSTOM_HOST_RUNTIME_VERSION, - code=rule.code, - family=rule.family, - slug=rule.slug, - message=rule.message, - remediation=rule.remediation, - subject=None if rule.subject is None else rule.subject.value, - implementation_fingerprint=fingerprint - if fingerprint is not None - else custom_rule_implementation_fingerprint( - rule=rule, - project_dir=project_dir, - import_closure=custom_rule_import_closure(rule=rule, project_dir=project_dir), - ), - config=config, - dialect=dialect, - project_dir=str(project_dir.resolve()), - ) - - -def _readset_digest(*, readset: tuple[FactKey, ...], digests: FactDigests) -> str | None: - try: - return digests.combined(readset) - except FactDigestError: - return None - - -def host_project(project: CompiledProjectView) -> CompiledProject: - """Drop project state no Rule fact reads, before digests and host payloads use it.""" - - return replace( - compiled_project_record(project, include_sql_expansions=False), - binding_catalog=None, - sql_expansions={}, - loaded_macros={}, - loader_functions=(), - hook_functions=(), - materialization_files=(), - external_sql_reference_resolver=None, - sql_scenarios=(), - sql_hook_files=(), - diagnostics=(), - unmatched_literal_sql_relations=(), - scope_index=ScopeIndex(), - ) - - -def _evaluations( - *, run: CustomHostRun, digests: FactDigests, cacheable: bool -) -> list[CustomEvaluationRow]: - """Each evaluated invocation with its read set and its digest when it may be cached.""" - - full_keys: tuple[FactKey, ...] | None = None - combined: dict[tuple[FactKey, ...], str | None] = {} - rows: list[CustomEvaluationRow] = [] - for (code, subject), result in run.evaluated.items(): - stored: tuple[tuple[FactKey, ...], str | None] | None = None - if cacheable and code not in run.uncacheable: - reads: tuple[FactKey, ...] | None = result.reads - if reads is None: - if full_keys is None: - full_keys = full_fact_keys(digests.views) - reads = (*run.untracked_reads.get(code, ()), *full_keys) - readset: tuple[FactKey, ...] = tuple(sorted(set(reads))) - if readset not in combined: - combined[readset] = _observed_digest( - readset=readset, observed=run.observed, digests=digests - ) - stored = (readset, combined[readset]) - rows.append((code, subject, stored, list(result.findings))) - return rows - - -def _observed_digest( - *, readset: tuple[FactKey, ...], observed: dict[FactKey, str | None], digests: FactDigests -) -> str | None: - """Digest a read set only when its filesystem facts still match what the rule observed.""" - - try: - for key in readset: - if is_tree_fact(key) and key in observed and observed[key] != digests.digest(key): - return None - return digests.combined(readset) - except FactDigestError: - return None - - -def _cpu_seconds() -> float: - """This thread's CPU plus that of finished child processes, which run the rule hosts.""" - - times: os.times_result = os.times() - return time.thread_time() + times.children_user + times.children_system - - -def _cpu_ms(started: float) -> int: - return int((_cpu_seconds() - started) * 1000) - - -def _elapsed_ms(started: float) -> int: - return round((time.monotonic() - started) * 1000) diff --git a/src/sqlbuild/rule_engine/_helpers/engine/rule_harness.py b/src/sqlbuild/rule_engine/_helpers/engine/rule_harness.py index 6717de9446..fb31d74cc8 100644 --- a/src/sqlbuild/rule_engine/_helpers/engine/rule_harness.py +++ b/src/sqlbuild/rule_engine/_helpers/engine/rule_harness.py @@ -6,7 +6,12 @@ from pathlib import Path from tempfile import TemporaryDirectory -from sqlbuild._native import build_compiled_project, discover_project_inputs, load_rules_config +from sqlbuild._native import ( + build_compiled_project, + discover_project_inputs, + evaluate_project_rules, + load_rules_config, +) from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter from sqlbuild.compiler.compile.types import CompiledProjectView from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs @@ -16,7 +21,6 @@ ) from sqlbuild.rule_engine.constants import MIN_CUSTOM_RULE_TEST_CASES from sqlbuild.rule_engine.exceptions import RuleAssertionError, RuleUsageError -from sqlbuild.rule_engine.main._evaluate import evaluate from sqlbuild.rule_engine.models import ( Rule, RuleCase, @@ -39,7 +43,7 @@ def run_rule_case(*, rule: RuleCheck | Rule, test_case: RuleCase) -> RuleResult: _ = _write_rule_project(root=root, rule=resolved, test_case=test_case) project: CompiledProjectView = _compile_rule_project(root=root) config: RulesConfig = load_rules_config(project_dir=root) - result: RulesResult = evaluate( + result: RulesResult = evaluate_project_rules( project=project, config=config, project_dir=root, diff --git a/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py b/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py deleted file mode 100644 index 3dbf3ea5d2..0000000000 --- a/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py +++ /dev/null @@ -1,443 +0,0 @@ -"""Run planned custom-rule invocations across isolated host processes and merge their results.""" - -from __future__ import annotations - -import os -import pickle -import shutil -import sys -import tempfile -import uuid -from collections.abc import Callable, Iterable -from concurrent.futures import FIRST_EXCEPTION, Future, ThreadPoolExecutor, wait -from dataclasses import replace -from pathlib import Path -from typing import Any, cast - -import orjson - -import sqlbuild._native as _native -from sqlbuild._native import ( - compiled_project_record, -) -from sqlbuild.compiler.compile.models import CompiledModel, CompiledProject -from sqlbuild.compiler.compile.types import CompiledModelView, CompiledProjectView -from sqlbuild.rule_engine._helpers.host.host_capacity import available_cores -from sqlbuild.rule_engine.constants import ( - CUSTOM_HOST_ABANDONED_SUFFIX, - CUSTOM_HOST_CANCELLED_MARKER, - CUSTOM_HOST_LAUNCH_MODULE, - CUSTOM_HOST_MAX_TRACKED_READS, - CUSTOM_HOST_RUNTIME_VERSION, - CUSTOM_RULE_PROJECT_SUBJECT, -) -from sqlbuild.rule_engine.exceptions import RulesError -from sqlbuild.rule_engine.models import ( - CustomHostEvaluation, - CustomHostRun, - CustomHostSlice, - Rule, - RulesConfig, -) -from sqlbuild.rule_engine.types import CustomHostPlan, FactKey, RuleSubject - -_HOST_TIMEOUT_MILLIS: int = 120_000 -_MIN_INVOCATIONS_PER_HOST: int = 64 -_MAX_HOSTS: int = 8 - - -def run_custom_hosts( - *, - project: CompiledProjectView, - config: RulesConfig, - project_dir: Path, - dialect: str, - rules: tuple[Rule, ...], - model_paths: tuple[str, ...], - plan: CustomHostPlan, - track_reads: bool, - verify_determinism: bool, -) -> CustomHostRun: - """Evaluate the plan across host processes and merge their results in single-host order.""" - - invocations: tuple[tuple[str, str], ...] = planned_invocations( - plan=plan, rules=rules, model_paths=model_paths - ) - plans: tuple[CustomHostPlan, ...] = ( - (plan,) - if verify_determinism - else partition_host_plan( - plan=plan, invocations=invocations, hosts=host_count(len(invocations)) - ) - ) - input_dir: Path = project_dir / "target" / "rules-cache" / "host-inputs" - input_dir.mkdir(parents=True, exist_ok=True) - input_path: Path = input_dir / f"project-{uuid.uuid4().hex}.pickle" - abandoned: Path = input_path.with_suffix(CUSTOM_HOST_ABANDONED_SUFFIX) - payload: dict[str, object] = { - "project_pickle_path": str(input_path.resolve()), - "abandoned_marker": str(abandoned.resolve()), - "parent_pid": os.getpid(), - "project_dir": str(project_dir.resolve()), - "dialect": dialect, - "verify_determinism": verify_determinism, - "track_reads": track_reads, - } - try: - responses: list[CustomHostSlice] = _run_partitions( - payload=payload, - plan=plan, - plans=plans, - input_dir=input_dir, - write_inputs=lambda: _write_inputs(path=input_path, project=project, config=config), - ) - finally: - input_path.unlink(missing_ok=True) - abandoned.unlink(missing_ok=True) - return _merge( - responses=responses, - order={invocation: index for index, invocation in enumerate(invocations)}, - ) - - -def _write_inputs(*, path: Path, project: CompiledProjectView, config: RulesConfig) -> None: - """Publish the project payload atomically; hosts started earlier wait for it to appear.""" - - payload: CompiledProjectView = host_payload_project(project) - with tempfile.NamedTemporaryFile( - dir=path.parent, prefix="writing-", suffix=".pickle", delete=False - ) as handle: - partial: Path = Path(handle.name) - try: - pickle.dump((payload, config), handle) - except BaseException: - handle.close() - partial.unlink(missing_ok=True) - raise - partial.replace(path) - - -def host_payload_project(project: CompiledProjectView) -> CompiledProject: - """Return the project hosts decode, without model state no Rule fact reads.""" - - return replace( - compiled_project_record(project), - models=tuple(map(_payload_model, project.models)), - native_project=None, - ) - - -def _payload_model(model: CompiledModelView) -> CompiledModelView: - """Drop model state no Rule fact reads, so hosts decode less of every model.""" - - return CompiledModel( - key=model.key, - deps=model.deps, - name=model.name, - relative_path=model.relative_path, - query_sql=model.query_sql, - config=model.config, - destination=model.destination, - schema_entry=model.schema_entry, - inferred_columns=model.inferred_columns, - fast_lineage_has_star=model.fast_lineage_has_star, - fast_lineage_star_resolved=model.fast_lineage_star_resolved, - authored_sql=model.authored_sql, - extract_implicit_alias_columns=model.extract_implicit_alias_columns, - enum_declarations=model.enum_declarations, - constant_declarations=model.constant_declarations, - binding_validated=model.binding_validated, - dynamic_column_contract=model.dynamic_column_contract, - rejected_sql_analysis_opt_out=model.rejected_sql_analysis_opt_out, - ) - - -def _abandon_inputs(marker: Path) -> bool: - """Tell waiting hosts no payload will come; report whether they can see the signal.""" - - try: - marker.parent.mkdir(parents=True, exist_ok=True) - marker.touch() - except OSError: - return False - return True - - -def _run_partitions( - *, - payload: dict[str, object], - plan: CustomHostPlan, - plans: tuple[CustomHostPlan, ...], - input_dir: Path, - write_inputs: Callable[[], None], -) -> list[CustomHostSlice]: - """Start hosts while the payload is written; failures and module state fall back to one host.""" - - if len(plans) == 1: - with ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-custom-host") as pool: - single: Future[CustomHostSlice] = pool.submit( - _run_host, payload={**payload, "plan": plan} - ) - _publish(write_inputs=write_inputs, payload=payload, futures=(single,)) - return [single.result()] - cancel_dir: Path = Path(tempfile.mkdtemp(dir=input_dir, prefix="hosts-")) - marker: Path = cancel_dir / CUSTOM_HOST_CANCELLED_MARKER - futures: list[Future[CustomHostSlice]] = [] - failed: bool = False - try: - with ThreadPoolExecutor( - max_workers=len(plans), thread_name_prefix="sqlbuild-custom-host" - ) as pool: - futures = [ - pool.submit( - _run_host, - payload={ - **payload, - "plan": partition, - "cancel_marker": str(marker.resolve()), - }, - ) - for partition in plans - ] - _publish(write_inputs=write_inputs, payload=payload, futures=futures) - done: set[Future[CustomHostSlice]] = wait(futures, return_when=FIRST_EXCEPTION).done - failed = any(future.exception() is not None for future in done) - if failed: - _cancel_hosts(marker) - finally: - shutil.rmtree(cancel_dir, ignore_errors=True) - if failed: - return [_run_host(payload={**payload, "plan": plan})] - slices: list[CustomHostSlice] = [future.result() for future in futures] - stateful: frozenset[str] = frozenset().union(*(item.stateful for item in slices)) - if not stateful: - return slices - rerun: CustomHostSlice = _run_host( - payload={ - **payload, - "plan": {code: plan[code] if code in stateful else [] for code in plan}, - }, - ) - return [*(_without_codes(response=item, codes=stateful) for item in slices), rerun] - - -def _publish( - *, - write_inputs: Callable[[], None], - payload: dict[str, object], - futures: Iterable[Future[CustomHostSlice]], -) -> None: - """Write the payload or release hosts; if release fails, the error waits for their deadline.""" - - try: - write_inputs() - except BaseException: - if _abandon_inputs(Path(str(payload["abandoned_marker"]))): - _ = wait(tuple(futures)) - raise - - -def _cancel_hosts(marker: Path) -> None: - """Stop hosts at their next invocation; a running invocation still ends or times out.""" - - marker.touch() - - -def _without_codes(*, response: CustomHostSlice, codes: frozenset[str]) -> CustomHostSlice: - """Drop every result a host reported for the given rules.""" - - return replace( - response, - evaluations=tuple(item for item in response.evaluations if item[0] not in codes), - untracked_reads={ - code: reads for code, reads in response.untracked_reads.items() if code not in codes - }, - uncacheable=response.uncacheable - codes, - stateful=response.stateful - codes, - ) - - -def planned_invocations( - *, plan: CustomHostPlan, rules: tuple[Rule, ...], model_paths: tuple[str, ...] -) -> tuple[tuple[str, str], ...]: - """List planned (rule, subject) invocations in the order one host evaluates them.""" - - planned: tuple[Rule, ...] = tuple( - sorted((rule for rule in rules if rule.code in plan), key=lambda rule: rule.code) - ) - invocations: list[tuple[str, str]] = [ - (rule.code, CUSTOM_RULE_PROJECT_SUBJECT) - for rule in planned - if rule.subject is RuleSubject.PROJECT - ] - model_rules: tuple[tuple[str, frozenset[str] | None], ...] = tuple( - (rule.code, None if plan[rule.code] is None else frozenset(plan[rule.code] or ())) - for rule in planned - if rule.subject is RuleSubject.MODEL - ) - for path in model_paths if model_rules else (): - invocations.extend( - (code, path) for code, selected in model_rules if selected is None or path in selected - ) - return tuple(invocations) - - -def partition_host_plan( - *, plan: CustomHostPlan, invocations: tuple[tuple[str, str], ...], hosts: int -) -> tuple[CustomHostPlan, ...]: - """Deal project invocations and whole model subjects round-robin across at most `hosts`.""" - - units: list[list[tuple[str, str]]] = [] - for code, subject in invocations: - if subject == CUSTOM_RULE_PROJECT_SUBJECT or not units or units[-1][0][1] != subject: - units.append([]) - units[-1].append((code, subject)) - count: int = max(1, min(hosts, len(units))) - if count == 1: - return (plan,) - partitions: list[dict[str, list[str]]] = [] - for _ in range(count): - partitions.append({code: [] for code in plan}) - for index, unit in enumerate(units): - for code, subject in unit: - partitions[index % count][code].append(subject) - return tuple(cast(CustomHostPlan, partition) for partition in partitions) - - -def host_count(invocations: int) -> int: - """Hosts worth starting: one per 64 invocations, bounded by usable CPUs and eight.""" - - return max(1, min(_MAX_HOSTS, available_cores(), invocations // _MIN_INVOCATIONS_PER_HOST)) - - -def _run_host(*, payload: dict[str, object]) -> CustomHostSlice: - try: - response_json: str = _native.run_custom_host_json( - orjson.dumps( - { - "program": sys.executable, - "arguments": ["-m", CUSTOM_HOST_LAUNCH_MODULE], - "timeout_millis": _HOST_TIMEOUT_MILLIS, - "runtime_version": CUSTOM_HOST_RUNTIME_VERSION, - "payload": payload, - } - ).decode() - ) - except (ValueError, TypeError) as error: - raise RulesError(str(error)) from error - return _decode_host_response(orjson.loads(response_json)) - - -def _decode_host_response(response: object) -> CustomHostSlice: - if not isinstance(response, dict): - raise RulesError("custom rule host returned an invalid payload") - payload: dict[str, Any] = cast(dict[str, Any], response) - raw_readsets: Any = payload.get("readsets") - raw_evaluations: Any = payload.get("evaluations") - raw_untracked: Any = payload.get("untracked") - raw_uncacheable: Any = payload.get("uncacheable") - raw_observed: Any = payload.get("observed") - if ( - not isinstance(raw_readsets, list) - or not isinstance(raw_evaluations, list) - or not isinstance(raw_untracked, dict) - or not isinstance(raw_uncacheable, list) - or not isinstance(raw_observed, list) - ): - raise RulesError("custom rule host returned an invalid payload") - readsets: list[tuple[FactKey, ...]] = list(map(decode_readset, raw_readsets)) - untracked: dict[str, tuple[FactKey, ...]] = { - str(code): decode_readset(reads) for code, reads in raw_untracked.items() - } - evaluations: list[tuple[str, str, CustomHostEvaluation]] = [] - for raw_item in raw_evaluations: - if not isinstance(raw_item, dict): - raise RulesError("custom rule host returned an invalid evaluation") - item: dict[str, Any] = cast(dict[str, Any], raw_item) - code: str = str(item["code"]) - subject: object = item.get("subject") - findings: Any = item.get("findings") - reads: object = item.get("reads") - if not isinstance(findings, list): - raise RulesError("custom rule host returned invalid findings") - tracked: bool = isinstance(reads, int) and code not in untracked - evaluations.append( - ( - code, - CUSTOM_RULE_PROJECT_SUBJECT if subject is None else str(subject), - CustomHostEvaluation( - findings=tuple(map(finding_payload, findings)), - reads=readsets[cast(int, reads)] if tracked else None, - ), - ) - ) - return CustomHostSlice( - evaluations=tuple(evaluations), - untracked_reads=untracked, - uncacheable=frozenset(str(code) for code in raw_uncacheable), - observed={ - tuple(map(str, key)): None if digest is None else str(digest) - for key, digest in raw_observed - }, - bounded=payload.get("bounded") is True, - stateful=frozenset(str(code) for code in payload.get("stateful") or ()), - ) - - -def _merge(*, responses: list[CustomHostSlice], order: dict[tuple[str, str], int]) -> CustomHostRun: - """Combine host slices exactly as one host would have reported every invocation.""" - - evaluations: list[tuple[str, str, CustomHostEvaluation]] = [] - untracked: set[str] = set() - for response in responses: - evaluations.extend(response.evaluations) - untracked.update(response.untracked_reads) - if len(responses) > 1: - evaluations.sort(key=lambda item: order.get((item[0], item[1]), len(order))) - tracked_readsets: set[tuple[FactKey, ...]] = { - result.reads - for code, _subject, result in evaluations - if result.reads is not None and code not in untracked - } - if any(response.bounded for response in responses) or ( - sum(map(len, tracked_readsets)) > CUSTOM_HOST_MAX_TRACKED_READS - ): - untracked.update(code for code, _subject, result in evaluations if result.reads is not None) - untracked_reads: dict[str, set[FactKey]] = {code: set() for code in untracked} - for response in responses: - for code, reads in response.untracked_reads.items(): - untracked_reads[code].update(reads) - evaluated: dict[tuple[str, str], CustomHostEvaluation] = {} - for code, subject, result in evaluations: - merged: CustomHostEvaluation = result - if code in untracked and result.reads is not None: - untracked_reads[code].update(result.reads) - merged = CustomHostEvaluation(findings=result.findings, reads=None) - evaluated[(code, subject)] = merged - return CustomHostRun( - evaluated=evaluated, - untracked_reads={code: tuple(sorted(reads)) for code, reads in untracked_reads.items()}, - uncacheable=frozenset().union(*(response.uncacheable for response in responses)), - observed=_merged_observations(responses), - ) - - -def _merged_observations(responses: list[CustomHostSlice]) -> dict[FactKey, str | None]: - """Keep each observed digest, forgetting any that two hosts observed differently.""" - - observed: dict[FactKey, str | None] = {} - for response in responses: - for key, digest in response.observed.items(): - if key in observed and observed[key] != digest: - observed[key] = None - else: - observed.setdefault(key, digest) - return observed - - -def decode_readset(value: Any) -> tuple[FactKey, ...]: - return tuple(tuple(map(str, key)) for key in value) - - -def finding_payload(value: Any) -> dict[str, object]: - return {str(key): item for key, item in value.items()} diff --git a/src/sqlbuild/rule_engine/_helpers/host/host_capacity.py b/src/sqlbuild/rule_engine/_helpers/host/host_capacity.py deleted file mode 100644 index 0f6d4e5840..0000000000 --- a/src/sqlbuild/rule_engine/_helpers/host/host_capacity.py +++ /dev/null @@ -1,103 +0,0 @@ -"""CPU capacity available to this process, including container CPU quotas.""" - -from __future__ import annotations - -import math -import os -from collections.abc import Callable, Iterator -from pathlib import Path - -_CGROUP_ROOT: Path = Path("/sys/fs/cgroup") -_PROC_CGROUP: Path = Path("/proc/self/cgroup") -_CGROUP_V1_CPU_CONTROLLER: str = "cpu" -_CGROUP_V1_CPU_DIRECTORIES: tuple[str, ...] = ("cpu", "cpu,cpuacct", "cpuacct,cpu") -_CPU_MAX_FIELD_COUNT: int = 2 - - -def available_cores(*, cgroup_root: Path = _CGROUP_ROOT, proc_cgroup: Path = _PROC_CGROUP) -> int: - """Return usable CPUs: the process CPU count capped by any cgroup v1 or v2 CPU quota.""" - - counter: Callable[[], int | None] | None = getattr(os, "process_cpu_count", None) - count: int = (counter() if counter is not None else _affinity_count()) or 1 - quotas: list[int] = list(_cgroup_quotas(cgroup_root=cgroup_root, proc_cgroup=proc_cgroup)) - return max(1, min((count, *quotas))) - - -def _affinity_count() -> int | None: - affinity: Callable[[int], set[int]] | None = getattr(os, "sched_getaffinity", None) - return len(affinity(0)) if affinity is not None else os.cpu_count() - - -def _cgroup_quotas(*, cgroup_root: Path, proc_cgroup: Path) -> Iterator[int]: - try: - entries: list[str] = proc_cgroup.read_text(encoding="utf-8").splitlines() - except (OSError, ValueError): - entries = [] - for entry in entries or ["0::/"]: - _hierarchy, _separator, remainder = entry.partition(":") - controllers, _separator, path = remainder.partition(":") - if not controllers: - yield from _quotas_in( - paths=tuple( - directory / "cpu.max" - for directory in _directories(base=cgroup_root, relative=path) - ), - parse=_cpu_max_quota, - ) - elif _CGROUP_V1_CPU_CONTROLLER in controllers.split(","): - for name in (controllers, *_CGROUP_V1_CPU_DIRECTORIES): - yield from _quotas_in( - paths=tuple(_directories(base=cgroup_root / name, relative=path)), - parse=_cfs_quota, - ) - - -def _directories(*, base: Path, relative: str) -> Iterator[Path]: - """Yield the cgroup directory and each ancestor down from the hierarchy root.""" - - directory: Path = base - yield directory - for part in Path(relative.strip("/")).parts: - directory = directory / part - yield directory - - -def _quotas_in(*, paths: tuple[Path, ...], parse: Callable[[Path], int | None]) -> Iterator[int]: - for path in paths: - quota: int | None = parse(path) - if quota is not None: - yield quota - - -def _cpu_max_quota(path: Path) -> int | None: - """Parse cgroup v2 `cpu.max`: `max ` means unlimited, else ` `.""" - - try: - fields: list[str] = path.read_text(encoding="utf-8").split() - except (OSError, ValueError): - return None - return ( - _quota(quota=fields[0], period=fields[1]) if len(fields) == _CPU_MAX_FIELD_COUNT else None - ) - - -def _cfs_quota(directory: Path) -> int | None: - """Parse cgroup v1 `cpu.cfs_quota_us` / `cpu.cfs_period_us`; a negative quota is unlimited.""" - - try: - quota: str = (directory / "cpu.cfs_quota_us").read_text(encoding="utf-8").strip() - period: str = (directory / "cpu.cfs_period_us").read_text(encoding="utf-8").strip() - except (OSError, ValueError): - return None - return _quota(quota=quota, period=period) - - -def _quota(*, quota: str, period: str) -> int | None: - try: - quota_us: int = int(quota) - period_us: int = int(period) - except ValueError: - return None - if quota_us <= 0 or period_us <= 0: - return None - return max(1, math.ceil(quota_us / period_us)) diff --git a/src/sqlbuild/rule_engine/_helpers/run/rule_threads.py b/src/sqlbuild/rule_engine/_helpers/run/rule_threads.py deleted file mode 100644 index 7b7d3ad7e4..0000000000 --- a/src/sqlbuild/rule_engine/_helpers/run/rule_threads.py +++ /dev/null @@ -1,148 +0,0 @@ -"""Thread and future handling around one rules run and early SQL lint.""" - -from __future__ import annotations - -from concurrent.futures import Executor, Future, ThreadPoolExecutor -from pathlib import Path -from threading import Event - -from sqlbuild._native import ( - build_lint_expansion_context, - early_lint_expansions, - finish_rule_run, - lint_stop_scope, - plan_rule_run, - run_lint, - run_planned_sql_rules, - sql_lint_preparation, -) -from sqlbuild.compiler.compile.models import ( - CompiledObjectKey, - CompiledSqlExpansion, - DeclarationScopeBuild, - SqlExpansionContext, -) -from sqlbuild.compiler.compile.types import CompileInputsView -from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs -from sqlbuild.compiler.pipeline.models import ProjectGraph -from sqlbuild.lint.models import LintConfig, LintRunResult -from sqlbuild.rule_engine._helpers.engine.custom_dispatch import evaluate_native, start_custom_rules -from sqlbuild.rule_engine.models import ( - CustomRulesOutcome, - PreparedSqlLint, - PreparedSqlLintResult, - RuleRunPlan, - RulesConfig, - RulesResult, - RulesRunResult, - SqlExpansionReuse, - SqlRulesEvaluation, -) - - -def evaluate_rules( - *, - graph: ProjectGraph, - discovered_inputs: DiscoveredProjectInputs, - config: RulesConfig, - project_dir: Path, - dialect: str, - selected_keys: frozenset[CompiledObjectKey] | None = None, - prepared_sql: PreparedSqlLint | None = None, - expansion_reuse: SqlExpansionReuse | None = None, - no_sql_analysis: bool = False, -) -> RulesRunResult: - """Evaluate independent rule phases concurrently, then finalize their combined findings.""" - - plan: RuleRunPlan = plan_rule_run( - graph=graph, - config=config, - project_dir=project_dir, - selected_keys=selected_keys, - no_sql_analysis=no_sql_analysis, - ) - with ( - ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-custom-rules") as custom, - ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-rules") as executor, - ): - custom_outcome: Future[CustomRulesOutcome] | None = start_custom_rules( - executor=custom, - project=plan.project, - config=plan.config, - project_dir=plan.project_dir, - catalogue=plan.catalogue, - custom_payloads=plan.custom_payloads, - dialect=dialect, - ) - native_result: Future[RulesResult] = executor.submit( - evaluate_native, - project=plan.project, - config=plan.config, - project_dir=plan.project_dir, - catalogue=plan.catalogue, - dialect=dialect, - defer_suppressions=True, - custom_payloads=plan.custom_payloads, - custom_outcome=custom_outcome, - ) - sql_result: SqlRulesEvaluation = run_planned_sql_rules( - plan=plan, - graph=graph, - discovered_inputs=discovered_inputs, - dialect=dialect, - prepared_sql=prepared_sql, - expansion_reuse=expansion_reuse, - ) - result: RulesResult = native_result.result() - return finish_rule_run(plan=plan, sql_result=sql_result, result=result) - - -def prepare_sql_rules( - *, inputs: CompileInputsView, executor: Executor, dialect: str, stop: Event -) -> PreparedSqlLint | None: - """Start early SQL lint on `executor` when the project selects SQL rules without a cache.""" - - preparation: tuple[Path, tuple[str, ...], LintConfig] | None = sql_lint_preparation( - inputs=inputs, dialect=dialect - ) - if preparation is None: - return None - project_dir, codes, config = preparation - return PreparedSqlLint( - codes=codes, - project_dir=project_dir, - future=executor.submit( - _prepare_sql_lint, - project_dir=project_dir, - config=config, - discovered_inputs=inputs.discovered_inputs, - static_declaration_scope=inputs.declaration_scope, - compiled_expansions=early_lint_expansions(inputs.model_inputs), - stop=stop, - ), - ) - - -def _prepare_sql_lint( - *, - project_dir: Path, - config: LintConfig, - discovered_inputs: DiscoveredProjectInputs, - static_declaration_scope: DeclarationScopeBuild | None, - compiled_expansions: dict[Path, CompiledSqlExpansion], - stop: Event, -) -> PreparedSqlLintResult: - with lint_stop_scope(stop): - context: SqlExpansionContext = build_lint_expansion_context( - project_dir=project_dir, - discovered_inputs=discovered_inputs, - static_declaration_scope=static_declaration_scope, - ) - result: LintRunResult = run_lint( - project_dir=project_dir, - config=config, - discovered_inputs=discovered_inputs, - compiled_expansions=compiled_expansions, - expansion_context=context, - ) - return PreparedSqlLintResult(result=result, context=context) diff --git a/src/sqlbuild/rule_engine/classes/fact_digests.py b/src/sqlbuild/rule_engine/classes/fact_digests.py deleted file mode 100644 index 521eae9dd6..0000000000 --- a/src/sqlbuild/rule_engine/classes/fact_digests.py +++ /dev/null @@ -1,52 +0,0 @@ -"""Memoized digests of recorded custom-rule facts.""" - -from __future__ import annotations - -import hashlib -from collections.abc import Callable, Iterable -from typing import Any - -import orjson - -from sqlbuild.rule_engine._helpers.engine.fact_replay import ( - fact_outcome_digest, - shared_fact_encoder, -) -from sqlbuild.rule_engine.classes.shared_fact_encoder import SharedFactEncoder -from sqlbuild.rule_engine.models import RuleFactViews -from sqlbuild.rule_engine.types import FactKey - - -class FactDigests: - """Memoized digests of recorded fact keys against one compiled project snapshot.""" - - def __init__(self, *, views: Callable[[], RuleFactViews]) -> None: - self._views_factory: Callable[[], RuleFactViews] = views - self._views: RuleFactViews | None = None - self._digests: dict[FactKey, str] = {} - self._shared: SharedFactEncoder = shared_fact_encoder() - - @property - def views(self) -> RuleFactViews: - if self._views is None: - self._views = self._views_factory() - return self._views - - def digest(self, key: FactKey) -> str: - """Return one fact digest, raising FactDigestError when it cannot be reproduced.""" - - cached: str | None = self._digests.get(key) - if cached is not None: - return cached - digest: str = fact_outcome_digest(views=self.views, key=key, shared=self._shared) - self._digests[key] = digest - return digest - - def combined(self, keys: Iterable[FactKey]) -> str: - """Digest one read set in canonical key order.""" - - combined: Any = hashlib.sha256() - for key in sorted(keys): - combined.update(orjson.dumps(key)) - combined.update(self.digest(key).encode()) - return combined.hexdigest() diff --git a/src/sqlbuild/rule_engine/main/_evaluate.py b/src/sqlbuild/rule_engine/main/_evaluate.py deleted file mode 100644 index 0677a2dd9d..0000000000 --- a/src/sqlbuild/rule_engine/main/_evaluate.py +++ /dev/null @@ -1,44 +0,0 @@ -"""Rules project evaluation entrypoint used by the public test harness.""" - -from pathlib import Path - -from sqlbuild._native import ( - build_rule_catalogue, - resolve_rule_ignore_selectors, - select_catalogue_rules, -) -from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.rule_engine._helpers.engine.custom_dispatch import evaluate_native -from sqlbuild.rule_engine._helpers.engine.hermeticity import verify_custom_rules -from sqlbuild.rule_engine.models import Rule, RulesConfig, RulesResult - - -def evaluate( - *, - project: CompiledProjectView, - config: RulesConfig, - project_dir: Path, - dialect: str = "generic", - verify_determinism: bool = False, -) -> RulesResult: - """Evaluate the selected compiler rules over compiled models.""" - - effective_config: RulesConfig = resolve_rule_ignore_selectors(config=config, project=project) - catalogue: tuple[Rule, ...] = build_rule_catalogue( - config=effective_config, project_dir=project_dir - ) - if any(rule.custom for rule in catalogue): - _ = verify_custom_rules( - rules=select_catalogue_rules( - catalogue=catalogue, config=effective_config, project_dir=project_dir - ), - project_dir=project_dir, - ) - return evaluate_native( - project=project, - config=effective_config, - project_dir=project_dir, - catalogue=catalogue, - dialect=dialect, - verify_determinism=verify_determinism, - ) diff --git a/src/sqlbuild/rule_engine/models.py b/src/sqlbuild/rule_engine/models.py index e13f8ce197..42c1434bbf 100644 --- a/src/sqlbuild/rule_engine/models.py +++ b/src/sqlbuild/rule_engine/models.py @@ -12,7 +12,7 @@ from sqlbuild.compiler.compile.types import CompiledProjectView from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs from sqlbuild.lint.models import LintRunResult -from sqlbuild.rule_engine.types import FactKey, RuleCheck, RuleOptionValue, RuleSubject +from sqlbuild.rule_engine.types import RuleCheck, RuleOptionValue, RuleSubject @dataclass(frozen=True, slots=True) @@ -467,36 +467,6 @@ class CustomRulesOutcome: custom_cpu_ms: int = 0 -@dataclass(frozen=True) -class CustomHostEvaluation: - """Decoded findings and recorded reads of one custom-rule invocation.""" - - findings: tuple[dict[str, object], ...] - reads: tuple[FactKey, ...] | None - - -@dataclass(frozen=True) -class CustomHostRun: - """Every planned custom-rule invocation, merged across host processes.""" - - evaluated: dict[tuple[str, str], CustomHostEvaluation] - untracked_reads: dict[str, tuple[FactKey, ...]] - uncacheable: frozenset[str] - observed: dict[FactKey, str | None] - - -@dataclass(frozen=True) -class CustomHostSlice: - """One host process's invocations, in the order that host evaluated them.""" - - evaluations: tuple[tuple[str, str, CustomHostEvaluation], ...] - untracked_reads: dict[str, tuple[FactKey, ...]] - uncacheable: frozenset[str] - observed: dict[FactKey, str | None] - bounded: bool - stateful: frozenset[str] - - @dataclass(frozen=True) class CustomHostPartition: """How one host of a split run detects module state and learns the run was cancelled.""" diff --git a/src/sqlbuild/rule_engine/types.py b/src/sqlbuild/rule_engine/types.py index b74529ea48..ccaedcbba2 100644 --- a/src/sqlbuild/rule_engine/types.py +++ b/src/sqlbuild/rule_engine/types.py @@ -26,13 +26,8 @@ type RuleOptionValue = bool | int | str | tuple[str, ...] | tuple[int, ...] type FactKey = tuple[str, ...] -type CustomEvaluationRow = tuple[ - str, str, tuple[tuple[FactKey, ...], str | None] | None, list[dict[str, object]] -] type CustomRulePlan = dict[str, frozenset[str] | None] -type CustomHostPlan = dict[str, list[str] | None] - class RuleSubject(StrEnum): """Compiler-owned unit used to invoke and cache one custom rule.""" diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/rules/helpers.py b/tests/e2e/src/sqlbuild/cli/commands/main/rules/helpers.py index 4f1330ff82..0a7ea9587d 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/rules/helpers.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/rules/helpers.py @@ -494,16 +494,20 @@ def few_marts(*, project: Project, ctx: RuleContext) -> list[Finding]: return [ctx.finding(subject=marts[0], message=f"{len(marts)} marts")] """, } -STALE_SOURCE_FACTS_SITECUSTOMIZE: str = """from sqlbuild.rule_engine.classes import fact_digests +STALE_SOURCE_FACTS_SITECUSTOMIZE: str = """from sqlbuild.rule_engine._helpers.engine import ( + fact_replay, +) -_digest = fact_digests.FactDigests.digest +_digest = fact_replay.fact_outcome_digest -def _stale_sources(self, key): - return "stale" if key == ("project.sources",) else _digest(self, key) +def _stale_sources(*, views, key, shared=None): + if key == ("project.sources",): + return "stale" + return _digest(views=views, key=key, shared=shared) -fact_digests.FactDigests.digest = _stale_sources +fact_replay.fact_outcome_digest = _stale_sources """ diff --git a/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py b/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py index 158d3ee9d7..007cdcfe63 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py @@ -13,8 +13,6 @@ import sqlbuild._native as native_module import sqlbuild.compiler.compile._helpers.render.user_macros as macros_module -import sqlbuild.rule_engine._helpers.host.custom_host_pool as custom_host_pool_module -import sqlbuild.rule_engine._helpers.run.rule_threads as rule_threads_module from sqlbuild._native import active_lint_stop, rules_bulk_cache_path from sqlbuild._native import run_cli as main from sqlbuild.compiler.compile.models import DeclarationScopeBuild, LoadedMacro @@ -22,7 +20,6 @@ from sqlbuild.compiler.sql_analysis.models import SqlLexicalSyntax from sqlbuild.lint.models import LintRunResult from tests.integration.src.sqlbuild.cli.commands.main._test_types import ( - CustomHostSplitIntegrationTestCase, CustomRuleCacheSequenceTestCase, DynamicPivotRulesIntegrationTestCase, EarlyLintInputsTestCase, @@ -1032,7 +1029,7 @@ def test_given_graph_error_when_compiling_then_early_lint_is_stopped_instead_of_ ) lint_started: Event = Event() stop_observed: list[bool] = [] - original_run_lint: Callable[..., LintRunResult] = rule_threads_module.run_lint + original_run_lint: Callable[..., LintRunResult] = native_module.run_lint def stop_awaiting_run_lint(**kwargs: Any) -> LintRunResult: lint_started.set() @@ -1040,7 +1037,7 @@ def stop_awaiting_run_lint(**kwargs: Any) -> LintRunResult: stop_observed.append(stop is not None and stop.wait(timeout=_EARLY_LINT_STOP_TIMEOUT_S)) return original_run_lint(**kwargs) - monkeypatch.setattr(rule_threads_module, "run_lint", stop_awaiting_run_lint) + monkeypatch.setattr(native_module, "run_lint", stop_awaiting_run_lint) exit_code: int = main(["--project-dir", str(tmp_path), "compile", "--json"]) @@ -1097,20 +1094,13 @@ def final_directory(*, model: Model, ctx: RuleContext) -> list[Finding]: """, encoding="utf-8", ) - host_runs: list[str] = [] - original_host: Callable[[str], str] = native_module.run_custom_host_json - - def recording_host(spec_json: str) -> str: - host_runs.append(spec_json) - return original_host(spec_json) - - monkeypatch.setattr(native_module, "run_custom_host_json", recording_host) - cold: set[str] = compile_finding_keys(project_dir=tmp_path, capsys=capsys) warm: set[str] = compile_finding_keys(project_dir=tmp_path, capsys=capsys) + _ = main(["--project-dir", str(tmp_path), "compile", "--json"]) + timings: dict[str, int] = json.loads(capsys.readouterr().out)["compile_timings"] - assert len(host_runs) == 1 + assert timings["rule_cache_misses"] == 0 assert warm == cold assert { f"models/orders.sql:{test_case.expected_code}", @@ -1118,80 +1108,6 @@ def recording_host(spec_json: str) -> str: } <= cold -@pytest.mark.parametrize( - "test_case", - [CustomHostSplitIntegrationTestCase("custom rules split across hosts", 9, 3, "XSQBRARCH001")], - ids=lambda case: case.description, -) -def test_given_custom_rules_split_across_hosts_when_compiling_then_output_matches_one_host( - test_case: CustomHostSplitIntegrationTestCase, - tmp_path: Path, - capsys: pytest.CaptureFixture[str], - monkeypatch: pytest.MonkeyPatch, -) -> None: - (tmp_path / "sqlbuild_project.toml").write_text( - 'name = "orders"\nadapter = "duckdb"\n\n' - f'[rules]\nselect = ["{test_case.expected_code}"]\n\n' - "[rules.thresholds]\nmin_custom_rule_test_cases = 0\n", - encoding="utf-8", - ) - for index in range(test_case.model_count): - folder: str = ("final", "staging", "staging")[index % 3] - model: Path = tmp_path / "models" / folder / f"orders_{index:02d}.sql" - model.parent.mkdir(parents=True, exist_ok=True) - model.write_text( - f'MODEL (description "Orders {index}");\nSELECT {index} AS order_id\n', - encoding="utf-8", - ) - rule_file: Path = tmp_path / "rules" / "architecture.py" - rule_file.parent.mkdir() - rule_file.write_text( - """from sqlbuild.rules import Finding, Model, RuleContext, rule - -@rule( - code="XSQBRARCH001", - message="Final models must use the final directory", - remediation="Move this model beneath models/final/.", -) -def final_directory(*, model: Model, ctx: RuleContext) -> list[Finding]: - return [] if "final" in model.path.parts else [ctx.finding(subject=model)] -""", - encoding="utf-8", - ) - host_runs: list[str] = [] - original_host: Callable[[str], str] = native_module.run_custom_host_json - - def recording_host(spec_json: str) -> str: - host_runs.append(spec_json) - return original_host(spec_json) - - def compile_payload(*, hosts: int) -> dict[str, Any]: - monkeypatch.setattr(custom_host_pool_module, "_MIN_INVOCATIONS_PER_HOST", 1) - monkeypatch.setattr(custom_host_pool_module, "available_cores", lambda: hosts) - exit_code: int = main(["--project-dir", str(tmp_path), "compile", "--json", "--no-cache"]) - assert exit_code == 1 - return json.loads(capsys.readouterr().out) - - monkeypatch.setattr(native_module, "run_custom_host_json", recording_host) - single: dict[str, Any] = compile_payload(hosts=1) - single_runs: int = len(host_runs) - rules_bulk_cache_path(project_dir=tmp_path, file_name="custom-rules.json").parent.parent.rename( - tmp_path / "single-rules-cache" - ) - split: dict[str, Any] = compile_payload(hosts=test_case.hosts) - split_runs: int = len(host_runs) - single_runs - warm: dict[str, Any] = compile_payload(hosts=test_case.hosts) - - assert (single_runs, split_runs, len(host_runs)) == (1, test_case.hosts, 1 + test_case.hosts) - assert split["diagnostics"] == single["diagnostics"] - assert len(split["diagnostics"]) == test_case.model_count - test_case.model_count // 3 - assert warm["diagnostics"] == single["diagnostics"] - assert warm["compile_timings"]["rule_cache_hits"] >= test_case.model_count - assert (tmp_path / "single-rules-cache" / "bulk" / "custom-rules.json").read_bytes() == ( - rules_bulk_cache_path(project_dir=tmp_path, file_name="custom-rules.json").read_bytes() - ) - - @pytest.mark.parametrize( "test_case", [RulesIntegrationTestCase("custom rule is enforced and cached", 1, "XSQBRARCH001")], diff --git a/tests/integration/src/sqlbuild/rule_engine/host_payload/_test_types.py b/tests/integration/src/sqlbuild/rule_engine/host_payload/_test_types.py deleted file mode 100644 index 63224ceda7..0000000000 --- a/tests/integration/src/sqlbuild/rule_engine/host_payload/_test_types.py +++ /dev/null @@ -1,8 +0,0 @@ -from dataclasses import dataclass - - -@dataclass(frozen=True) -class HostPayloadFactsTestCase: - description: str - project_files: dict[str, str] - expected_unpopulated_trimmed_fields: frozenset[str] diff --git a/tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py b/tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py deleted file mode 100644 index f31880b78a..0000000000 --- a/tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py +++ /dev/null @@ -1,132 +0,0 @@ -import dataclasses -from itertools import product -from pathlib import Path -from typing import Any - -import pytest - -from sqlbuild._native import build_compiled_project, discover_project_inputs -from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter -from sqlbuild.compiler.compile.models import CompiledModel, CompiledProject -from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.rule_engine._helpers.engine import custom_rules, fact_replay -from sqlbuild.rule_engine._helpers.engine.fact_replay import ( - fact_outcome_digest, - shared_fact_encoder, -) -from sqlbuild.rule_engine._helpers.host import custom_host_pool -from sqlbuild.rule_engine.classes.project_tree import public_model -from sqlbuild.rule_engine.classes.rule_context import RuleFactViews, build_rule_fact_views -from sqlbuild.rule_engine.models import Model -from sqlbuild.rule_engine.types import FactKey - -_TREE_ARGUMENTS: tuple[str, ...] = ("models/marts/order_totals.sql", "models") - - -class OrdersReferenceResolver: - """External reference resolver that knows no extra relations.""" - - def validate_model_names(self, *, known_model_names: set[str]) -> None: - del known_model_names - - def extend_sql_test_model_names(self, *, known_model_names: set[str]) -> set[str]: - return known_model_names - - def extend_sql_test_source_names(self, *, known_source_names: set[str]) -> set[str]: - return known_source_names - - def extend_sql_test_seed_names(self, *, known_seed_names: set[str]) -> set[str]: - return known_seed_names - - def validate_reference( - self, - *, - ref_kind: str, - ref_name: str, - ref_package: str | None, - owner_relative_sql_path: Path, - ) -> None: - del ref_kind, ref_name, ref_package, owner_relative_sql_path - - def resolve_reference( - self, *, ref_kind: str, ref_name: str, ref_package: str | None - ) -> str | None: - del ref_kind, ref_name, ref_package - return None - - -def compile_files(*, root: Path, files: dict[str, str]) -> CompiledProjectView: - """Write and compile one DuckDB project with an external reference resolver.""" - - for relative_path, contents in files.items(): - path: Path = root / relative_path - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(contents, encoding="utf-8") - return build_compiled_project( - discovered_inputs=discover_project_inputs(project_dir=root), - adapter=DuckDbAdapter(), - external_sql_reference_resolver=OrdersReferenceResolver(), - ) - - -def dropped_fields(*, before: object, after: object) -> frozenset[str]: - """Return the fields of `after`'s dataclass whose values differ between the two snapshots.""" - - names: tuple[str, ...] = tuple(field.name for field in dataclasses.fields(type(after))) - return frozenset(filter(lambda name: getattr(before, name) != getattr(after, name), names)) - - -def every_fact_key(project: CompiledProjectView) -> tuple[FactKey, ...]: - """Return a key for every fact Rules can read, over every model of the project.""" - - models: tuple[Model, ...] = tuple(public_model(model) for model in project.models) - return ( - *((fact,) for fact in fact_replay._GLOBAL_FACTS), - *( - (fact, model.path.as_posix()) - for fact, model in product(fact_replay._MODEL_PATH_FACTS, models) - ), - *((fact, model.name) for fact, model in product(fact_replay._MODEL_NAME_FACTS, models)), - *( - (fact, *_TREE_ARGUMENTS[: fact_replay._TEXT_FACT_ARITY[fact]]) - for fact in fact_replay._TEXT_FACTS - ), - ) - - -def fact_digests(*, project: CompiledProjectView, root: Path) -> dict[FactKey, tuple[str, str]]: - """Digest every fact as hosts observe it and as whole-project digests encode it.""" - - views: RuleFactViews = build_rule_fact_views( - project=project, project_dir=root, dialect="duckdb" - ) - return { - key: ( - fact_outcome_digest(views=views, key=key), - fact_outcome_digest(views=views, key=key, shared=shared_fact_encoder()), - ) - for key in every_fact_key(project) - } - - -def record_trimmed_fields( - monkeypatch: pytest.MonkeyPatch, -) -> dict[type[object], set[str]]: - """Record replaced fields and omitted constructor fields in the host-payload trimmers.""" - - trimmed: dict[type[object], set[str]] = {CompiledProject: set(), CompiledModel: set()} - - def recording_replace(obj: Any, /, **changes: object) -> Any: - trimmed.setdefault(type(obj), set()).update(changes) - return dataclasses.replace(obj, **changes) - - def recording_model(**values: Any) -> CompiledModel: - trimmed[CompiledModel].update( - {field.name for field in dataclasses.fields(CompiledModel)} - values.keys() - ) - return CompiledModel(**values) - - monkeypatch.setattr(custom_rules, "replace", recording_replace) - monkeypatch.setattr(custom_host_pool, "replace", recording_replace) - monkeypatch.setattr(custom_host_pool, "CompiledModel", recording_model) - return trimmed diff --git a/tests/integration/src/sqlbuild/rule_engine/host_payload/test_host_payload_facts.py b/tests/integration/src/sqlbuild/rule_engine/host_payload/test_host_payload_facts.py deleted file mode 100644 index b6ad426c3f..0000000000 --- a/tests/integration/src/sqlbuild/rule_engine/host_payload/test_host_payload_facts.py +++ /dev/null @@ -1,136 +0,0 @@ -"""Every Rule fact reads the same values from the trimmed host payload as from the full project.""" - -from pathlib import Path - -import pytest - -from sqlbuild._native import ( - compiled_project_record, -) -from sqlbuild.compiler.compile.models import CompiledModel, CompiledProject -from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.rule_engine._helpers.engine.custom_rules import host_project -from sqlbuild.rule_engine._helpers.host.custom_host_pool import host_payload_project -from sqlbuild.rule_engine.classes.project_tree import public_model -from sqlbuild.rule_engine.types import FactKey -from tests.integration.src.sqlbuild.rule_engine.host_payload._test_types import ( - HostPayloadFactsTestCase, -) -from tests.integration.src.sqlbuild.rule_engine.host_payload.helpers import ( - compile_files, - dropped_fields, - fact_digests, - record_trimmed_fields, -) - - -@pytest.mark.parametrize( - "test_case", - ( - HostPayloadFactsTestCase( - description="sources, refs, contracts, macros, declarations, tests, and audits", - project_files={ - "sqlbuild_project.toml": 'name = "orders"\nadapter = "duckdb"\n', - "sources/raw.yml": ( - "sources:\n" - " - name: raw_orders\n description: Raw orders.\n" - ' expression: "(SELECT 1 AS id, CAST(10.5 AS DOUBLE) AS amount)"\n' - " columns:\n - name: id\n type: INTEGER\n" - " - name: raw_customers\n description: Raw customers.\n" - ' expression: "(SELECT 1 AS customer_id)"\n' - ), - "constants/policy.sql": "CONSTANT (name minimum_amount, value 7);\n", - "enums/order_status.sql": "ENUM (name order_status, members [PAID, OPEN]);\n", - "models/marts/_sqlbuild/_macros/amounts.py": ( - 'def doubled(amount: str) -> str:\n return f"({amount}) * 2"\n' - ), - "models/staging/stg_orders.sql": ( - "MODEL (description 'Staged orders.',\n contract enforced,\n" - " columns (\n" - " order_id (type INTEGER, nullable false, audits [not_null]),\n" - " amount (type DOUBLE),\n status (type order_status),\n ),\n);\n\n" - "SELECT CAST(id AS INTEGER) AS order_id, amount, 'PAID' AS status\n" - 'FROM __source("raw_orders")\n' - ), - "models/marts/order_totals.sql": ( - 'MODEL (description "Order totals", materialized table);\n\n' - 'SELECT o.order_id, @doubled("o.amount") AS doubled_amount, o.status\n' - 'FROM __ref("stg_orders") AS o WHERE o.amount > 7\n' - ), - "models/marts/order_flags.sql": ( - 'MODEL (description "Order flags");\n\n' - 'SELECT o.order_id, o.missing_flag AS flag FROM __ref("stg_orders") AS o\n' - ), - "models/marts/flag_counts.sql": ( - 'MODEL (description "Flag counts");\n\n' - 'SELECT f.flag FROM __ref("order_flags") AS f\n' - ), - "hooks/sql/analyze_orders.sql": ( - 'HOOK (description "Analyze orders");\n\nSELECT 1\n' - ), - "hooks/python/order_hooks.py": ( - "from sqlbuild.hooks import hook\n\n" - "@hook\n" - "def log_orders(ctx):\n" - " '''Log orders.'''\n" - " ctx.execute_sql('INSERT INTO audit_log SELECT 1')\n" - ), - "python/loaders/order_rows.py": ( - "from sqlbuild.loaders import loader\n\n" - "@loader(columns=[{'name': 'order_id', 'type': 'INTEGER'}])\n" - "def fetch_orders(ctx):\n" - " '''Fetch orders.'''\n" - " return [{'order_id': 1}]\n" - ), - "materializations/copy_table.py": ( - "from sqlbuild.executor.custom.models import (\n" - " MaterializationContext,\n MaterializationResult,\n)\n\n\n" - "def materialize(ctx: MaterializationContext) -> MaterializationResult:\n" - ' ctx.execute_sql(f"CREATE TABLE {ctx.destination} AS {ctx.sql}")\n' - " return MaterializationResult(relation=ctx.destination)\n" - ), - "tests/scenarios/paid_orders.sql": ( - "SCENARIO (description 'Paid orders.');\n\n" - "WITH\n__source__raw_orders AS (SELECT 1 AS id, 10.5 AS amount),\n" - "__expected__stg_orders AS (\n" - " SELECT 1 AS order_id, 10.5 AS amount, 'PAID' AS status\n)\n" - "SELECT 1\n" - ), - "tests/unit/test_order_totals.sql": ( - "TEST ();\n\nWITH\n__ref__stg_orders AS (\n" - " SELECT 1 AS order_id, 10.0 AS amount\n),\n" - "__expected__order_totals AS (\n" - " SELECT 1 AS order_id, 20.0 AS doubled_amount, 'PAID' AS status\n)\n" - "SELECT 1\n" - ), - }, - expected_unpopulated_trimmed_fields=frozenset(), - ), - ), - ids=lambda case: case.description, -) -def test_given_compiled_project_when_trimming_host_payload_then_every_fact_is_unchanged( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostPayloadFactsTestCase -) -> None: - project: CompiledProjectView = compile_files(root=tmp_path, files=test_case.project_files) - trimmed_fields: dict[type[object], set[str]] = record_trimmed_fields(monkeypatch) - trimmed: CompiledProjectView = host_payload_project(host_project(project)) - - full_digests: dict[FactKey, tuple[str, str]] = fact_digests(project=project, root=tmp_path) - trimmed_digests: dict[FactKey, tuple[str, str]] = fact_digests(project=trimmed, root=tmp_path) - - assert trimmed_digests == full_digests - assert tuple(map(public_model, trimmed.models)) == tuple(map(public_model, project.models)) - assert ( - dropped_fields(before=compiled_project_record(project), after=trimmed) - == trimmed_fields[CompiledProject] - test_case.expected_unpopulated_trimmed_fields - ) - assert ( - frozenset().union( - *( - dropped_fields(before=model, after=trimmed_model) - for model, trimmed_model in zip(project.models, trimmed.models, strict=True) - ) - ) - == trimmed_fields[CompiledModel] - test_case.expected_unpopulated_trimmed_fields - ) diff --git a/tests/unit/src/sqlbuild/compiler/pipeline/main/test_compile_phase_timings.py b/tests/unit/src/sqlbuild/compiler/pipeline/main/test_compile_phase_timings.py index f580a17fd0..e5dba7c66c 100644 --- a/tests/unit/src/sqlbuild/compiler/pipeline/main/test_compile_phase_timings.py +++ b/tests/unit/src/sqlbuild/compiler/pipeline/main/test_compile_phase_timings.py @@ -15,6 +15,7 @@ from sqlbuild.compiler.pipeline._helpers import materializations from sqlbuild.compiler.pipeline.models import CompilePipelineOptions, CompilePipelineResult from sqlbuild.compiler.planner.models import PlanOutput +from sqlbuild.diagnostics.models import PartialBuildPhaseTimings from sqlbuild.rule_engine._helpers.run import rule_threads from sqlbuild.rule_engine.models import RulesRunResult from sqlbuild.runtime.contracts.models import ConnectionHooks @@ -80,7 +81,7 @@ def test_given_slow_planning_connection_when_compiling_then_direct_phases_are_di (tmp_path / "sqlbuild_project.toml").write_text('name = "timing"\n', encoding="utf-8") monkeypatch.setattr(_native, "resolve_configured_rule_selection", Mock(return_value=None)) monkeypatch.setattr( - rule_threads, + _native, "evaluate_rules", Mock( return_value=RulesRunResult(findings=(), evaluated_models=0, built_in_ms=0, custom_ms=0) diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/helpers.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/helpers.py index 3c79ba9238..535c17cbe0 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/helpers.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/helpers.py @@ -9,6 +9,7 @@ import pytest +import sqlbuild._native as native from sqlbuild import _native from sqlbuild._native import ( NativeRulesRequest, @@ -16,6 +17,8 @@ install_rules_skills, project_with, ) +from sqlbuild._native import evaluate_custom_rules as evaluate_custom_rules_cached +from sqlbuild._native import evaluate_project_rules as evaluate from sqlbuild.compiler.compile.models import ( CompiledDirectLogicSqlTestPayload, CompiledModelSqlTestPayload, @@ -29,7 +32,7 @@ from sqlbuild.compiler.scopes.models import ScopeIndex from sqlbuild.rule_engine._helpers.engine import custom_dispatch as native from sqlbuild.rule_engine._helpers.engine.custom_rules import evaluate_custom_rules_cached -from sqlbuild.rule_engine.constants import MIN_CUSTOM_RULE_TEST_CASES +from sqlbuild.rule_engine.constants import MIN_CUSTOM_RULE_TEST_CASES, RULES_SKILL_PATHS from sqlbuild.rule_engine.main._evaluate import evaluate from sqlbuild.rule_engine.models import ( CustomRulesOutcome, @@ -101,7 +104,7 @@ def evaluate_rules_request(request: object) -> NativeRulesEvaluation: monkeypatch.setattr(_native, "scope_metadata_projection", scope_metadata_projection) monkeypatch.setattr(native, "build_rules_request", build_rules_request) monkeypatch.setattr(native, "evaluate_rules_request", evaluate_rules_request) - native.evaluate_native( + native.evaluate_native_rules( project=project, config=RulesConfig(cache=RulesCacheConfig(enabled=False)), project_dir=project_dir, @@ -383,7 +386,7 @@ def evaluate_incremental_step( rule_exceptions=step.rule_exceptions, cache=RulesCacheConfig(enabled=cache_enabled), ) - return native.evaluate_native( + return native.evaluate_native_rules( project=_incremental_project(step.models), config=config, project_dir=project_dir, diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_catalogue.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_catalogue.py index 2bce72eb50..292607f670 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_catalogue.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_catalogue.py @@ -10,7 +10,7 @@ install_rules_skills, select_catalogue_rules, ) -from sqlbuild.rule_engine.main._evaluate import evaluate +from sqlbuild._native import evaluate_project_rules as evaluate from sqlbuild.rule_engine.models import ( Rule, RulesCacheConfig, diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rule_evidence.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rule_evidence.py index 89517c73b8..e8bd7e90fe 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rule_evidence.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rule_evidence.py @@ -9,9 +9,9 @@ custom_rule_implementation_fingerprint, custom_rule_test_evidence, ) +from sqlbuild._native import evaluate_project_rules as evaluate from sqlbuild.compiler.compile.types import CompiledProjectView from sqlbuild.rule_engine.exceptions import RulesError -from sqlbuild.rule_engine.main._evaluate import evaluate from sqlbuild.rule_engine.models import Rule, RulesConfig, RulesResult from tests.unit.src.sqlbuild.rule_engine._helpers.engine._test_types import ( CustomRuleEvidenceTestCase, diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rule_reads.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rule_reads.py index 08dae222f7..ca53f4685a 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rule_reads.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rule_reads.py @@ -8,7 +8,7 @@ import sqlbuild._native as native_module from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.rule_engine._helpers.engine import custom_rules +from sqlbuild.rule_engine._helpers.engine import fact_replay from sqlbuild.rule_engine.models import CustomRulesOutcome from tests.unit.src.sqlbuild.rule_engine._helpers.engine._test_types import ( CustomRuleProjectFileTestCase, @@ -253,14 +253,13 @@ def test_given_file_edited_while_rules_run_when_rerunning_then_new_contents_are_ project: CompiledProjectView = two_model_project( customers_sql=_CUSTOMERS_SQL, customers_config={} ) - run_host: Callable[..., object] = custom_rules.run_custom_hosts + digest: Callable[..., str] = fact_replay.fact_outcome_digest - def run_host_then_edit(**kwargs: Any) -> object: - result: object = run_host(**kwargs) - owners.write_text(_OWNERS_AFTER, encoding="utf-8") - return result + def edit_then_digest(**kwargs: Any) -> str: + _ = owners.write_text(_OWNERS_AFTER, encoding="utf-8") + return digest(**kwargs) - monkeypatch.setattr(custom_rules, "run_custom_hosts", run_host_then_edit) + monkeypatch.setattr(fact_replay, "fact_outcome_digest", edit_then_digest) during_edit: CustomRulesOutcome = evaluate_cached_custom_rules( project=project, project_dir=tmp_path, cache_enabled=True ) diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rules.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rules.py index 43d9d99bfa..9f1ea679d4 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rules.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rules.py @@ -9,10 +9,10 @@ compiled_model_record, project_with, ) +from sqlbuild._native import evaluate_project_rules as evaluate from sqlbuild.compiler.compile.models import CompactLineageFacts from sqlbuild.compiler.compile.types import CompiledProjectView from sqlbuild.rule_engine.exceptions import RulesError -from sqlbuild.rule_engine.main._evaluate import evaluate from sqlbuild.rule_engine.models import ( RuleExemption, RuleIgnore, diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py index 9c581503f4..5f235cc554 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py @@ -2,7 +2,6 @@ import os import random -import shutil import time from collections.abc import Callable from operator import attrgetter @@ -12,14 +11,12 @@ import pytest -import sqlbuild._native as native_module from sqlbuild._native import ( build_rule_catalogue, project_with, ) +from sqlbuild._native import evaluate_custom_rules as evaluate_custom_rules_cached from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.rule_engine._helpers.engine.custom_rules import evaluate_custom_rules_cached -from sqlbuild.rule_engine._helpers.host import custom_host_pool from sqlbuild.rule_engine._helpers.host.module_state import module_state_token, rule_namespaces from sqlbuild.rule_engine.classes.module_state_monitor import ModuleStateMonitor from sqlbuild.rule_engine.exceptions import OpaqueModuleStateError @@ -112,34 +109,6 @@ def evaluate_order_rules( ) -def use_hosts(*, monkeypatch: pytest.MonkeyPatch, hosts: int) -> list[str]: - """Allow up to `hosts` host processes for any plan and record every host launch.""" - - launches: list[str] = [] - launch: Callable[[str], str] = native_module.run_custom_host_json - - def recording_launch(spec_json: str) -> str: - launches.append(spec_json) - return launch(spec_json) - - monkeypatch.setattr(custom_host_pool, "_MIN_INVOCATIONS_PER_HOST", 1) - monkeypatch.setattr(custom_host_pool, "available_cores", lambda: hosts) - monkeypatch.setattr(native_module, "run_custom_host_json", recording_launch) - return launches - - -def write_cgroup_files(*, root: Path, proc_cgroup: str, files: dict[str, str]) -> Path: - """Write a fake `/proc/self/cgroup` and cgroup hierarchy beneath `root`.""" - - for relative_path, contents in files.items(): - path: Path = root / "cgroup" / relative_path - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(contents, encoding="utf-8", errors="surrogateescape") - proc: Path = root / "proc-self-cgroup" - proc.write_text(proc_cgroup, encoding="utf-8", errors="surrogateescape") - return proc - - def record_signals(*, monkeypatch: pytest.MonkeyPatch) -> list[tuple[int, int]]: """Replace process signalling with a recorder so no process can be signalled.""" @@ -442,47 +411,6 @@ def timed_state_history( return time.process_time() - started, history -def delay_inputs(*, monkeypatch: pytest.MonkeyPatch, delay_seconds: float) -> None: - """Publish every host project payload only after a delay.""" - - write: Callable[..., None] = custom_host_pool._write_inputs - - def delayed_write(*, path: Path, project: CompiledProjectView, config: RulesConfig) -> None: - time.sleep(delay_seconds) - write(path=path, project=project, config=config) - - monkeypatch.setattr(custom_host_pool, "_write_inputs", delayed_write) - - -def fail_inputs(*, monkeypatch: pytest.MonkeyPatch, delay_seconds: float) -> None: - """Fail to publish every host project payload after a delay.""" - - def failed_write(*, path: Path, project: CompiledProjectView, config: RulesConfig) -> None: - del path, project, config - time.sleep(delay_seconds) - raise OSError("project payload disk is full") - - monkeypatch.setattr(custom_host_pool, "_write_inputs", failed_write) - - -def fail_inputs_without_directory( - *, monkeypatch: pytest.MonkeyPatch, delay_seconds: float -) -> list[float]: - """Remove the host-input directory, then fail to publish; return when each write failed.""" - - failures: list[float] = [] - - def failed_write(*, path: Path, project: CompiledProjectView, config: RulesConfig) -> None: - del project, config - time.sleep(delay_seconds) - shutil.rmtree(path.parent) - failures.append(time.monotonic()) - raise OSError("project payload disk is full") - - monkeypatch.setattr(custom_host_pool, "_write_inputs", failed_write) - return failures - - def touch_files(*, root: Path, names: tuple[str, ...]) -> None: """Create empty files beneath one folder.""" diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py index d14fe28676..54a2b919f1 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py @@ -1,408 +1,17 @@ -"""Custom rules split across host processes report exactly what one host reports.""" +"""The custom-rule host waits for its payload only while the parent still wants it.""" import time from pathlib import Path import pytest -from sqlbuild._native import rules_bulk_cache_path -from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.rule_engine._helpers.host import custom_host, custom_host_pool -from sqlbuild.rule_engine.constants import CUSTOM_RULES_CACHE_FILE +from sqlbuild.rule_engine._helpers.host import custom_host from sqlbuild.rule_engine.exceptions import HostCancelledError, RulesError -from sqlbuild.rule_engine.models import CustomRulesOutcome from tests.unit.src.sqlbuild.rule_engine._helpers.host._test_types import ( AwaitInputsTestCase, AwaitPublishedInputsTestCase, - HostCancellationTestCase, - HostFailureTestCase, - HostInputsTestCase, - HostPartitionTestCase, - HostPlanTestCase, ) -from tests.unit.src.sqlbuild.rule_engine._helpers.host.helpers import ( - delay_inputs, - evaluate_order_rules, - fail_inputs, - fail_inputs_without_directory, - orders_project, - record_signals, - touch_files, - use_hosts, - write_order_rules, -) - - -@pytest.mark.parametrize( - "test_case", - [ - HostPartitionTestCase( - description="three hosts", model_count=7, hosts=3, expected_host_runs=3 - ), - HostPartitionTestCase( - description="more hosts than subjects", model_count=2, hosts=8, expected_host_runs=3 - ), - HostPartitionTestCase( - description="module state reruns the rule in one host", - model_count=7, - hosts=3, - expected_host_runs=4, - records_module_state=True, - ), - HostPartitionTestCase( - description="cross-model duplicate detector reruns in one host", - model_count=7, - hosts=3, - expected_host_runs=4, - detects_duplicates=True, - ), - HostPartitionTestCase( - description="uncached cross-model duplicate detector reruns in one host", - model_count=7, - hosts=3, - expected_host_runs=4, - detects_duplicates=True, - cache_enabled=False, - ), - ], - ids=lambda case: case.description, -) -def test_given_partitioned_hosts_when_evaluating_then_findings_and_cache_match_one_host( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostPartitionTestCase -) -> None: - write_order_rules( - root=tmp_path, - records_module_state=test_case.records_module_state, - detects_duplicates=test_case.detects_duplicates, - ) - project: CompiledProjectView = orders_project(model_count=test_case.model_count) - cache_path: Path = rules_bulk_cache_path( - project_dir=tmp_path, file_name=CUSTOM_RULES_CACHE_FILE - ) - cache_glob: str = cache_path.relative_to(tmp_path).as_posix() - cache_enabled: bool = test_case.cache_enabled - - single_launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=1) - single: CustomRulesOutcome = evaluate_order_rules( - project=project, root=tmp_path, cache_enabled=cache_enabled - ) - single_warm: CustomRulesOutcome = evaluate_order_rules( - project=project, root=tmp_path, cache_enabled=cache_enabled - ) - single_cache: bytes = b"".join(path.read_bytes() for path in tmp_path.glob(cache_glob)) - cache_path.unlink(missing_ok=True) - monkeypatch.undo() - split_launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) - split: CustomRulesOutcome = evaluate_order_rules( - project=project, root=tmp_path, cache_enabled=cache_enabled - ) - split_cold_launches: int = len(split_launches) - split_cache: bytes = b"".join(path.read_bytes() for path in tmp_path.glob(cache_glob)) - split_warm: CustomRulesOutcome = evaluate_order_rules( - project=project, root=tmp_path, cache_enabled=cache_enabled - ) - - assert split.findings == single.findings - assert split_warm.findings == single_warm.findings == single.findings - assert len(single_launches) == 1 + int(not cache_enabled) - assert split_cold_launches == test_case.expected_host_runs - assert (split.cache_hits, split.cache_misses) == (single.cache_hits, single.cache_misses) - assert (split_warm.cache_hits, split_warm.cache_misses) == ( - single_warm.cache_hits, - single_warm.cache_misses, - ) - assert split_cache == single_cache - assert bool(single_cache) is cache_enabled - assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) - - -@pytest.mark.parametrize( - "test_case", - [ - HostPartitionTestCase( - description="edit after split", model_count=6, hosts=3, expected_host_runs=3 - ) - ], - ids=lambda case: case.description, -) -def test_given_results_cached_by_split_hosts_when_rerunning_then_only_edited_subjects_rerun( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostPartitionTestCase -) -> None: - write_order_rules(root=tmp_path) - project: CompiledProjectView = orders_project(model_count=test_case.model_count) - edited: CompiledProjectView = orders_project( - model_count=test_case.model_count, edits={"orders_003": "SELECT 2 AS order_id"} - ) - launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) - - cold: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) - cold_launches: int = len(launches) - warm: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) - warm_launches: int = len(launches) - cold_launches - after_edit: CustomRulesOutcome = evaluate_order_rules(project=edited, root=tmp_path) - oracle: CustomRulesOutcome = evaluate_order_rules( - project=edited, root=tmp_path, cache_enabled=False - ) - - assert cold_launches == test_case.expected_host_runs - assert warm_launches == 0 - assert (warm.cache_hits, warm.cache_misses) == (cold.cache_misses, 0) - assert warm.findings == cold.findings - assert (after_edit.cache_hits, after_edit.cache_misses) == (cold.cache_misses - 1, 1) - assert after_edit.findings == oracle.findings - - -@pytest.mark.parametrize( - "test_case", - [ - HostFailureTestCase( - description="earliest failure lives on a later host", - model_count=6, - failing_models=("orders_001", "orders_003"), - expected_failing_model="orders_001", - ), - HostFailureTestCase( - description="one failing subject", - model_count=6, - failing_models=("orders_005",), - expected_failing_model="orders_005", - ), - ], - ids=lambda case: case.description, -) -def test_given_failing_rule_on_split_hosts_when_evaluating_then_single_host_error_is_raised( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostFailureTestCase -) -> None: - write_order_rules(root=tmp_path, failing=test_case.failing_models) - project: CompiledProjectView = orders_project(model_count=test_case.model_count) - - _ = use_hosts(monkeypatch=monkeypatch, hosts=1) - with pytest.raises(RulesError) as single: - _ = evaluate_order_rules(project=project, root=tmp_path) - monkeypatch.undo() - _ = use_hosts(monkeypatch=monkeypatch, hosts=3) - with pytest.raises(RulesError) as split: - _ = evaluate_order_rules(project=project, root=tmp_path) - - assert str(split.value) == str(single.value) - assert f"cannot inspect {test_case.expected_failing_model}" in str(split.value) - - -@pytest.mark.parametrize( - "test_case", - [ - HostPlanTestCase( - description="project rules and whole subjects are dealt round-robin", - plan={"XSQBRT101": None, "XSQBRT102": None, "XSQBRT103": ["models/b.sql"]}, - invocations=( - ("XSQBRT102", ""), - ("XSQBRT101", "models/a.sql"), - ("XSQBRT101", "models/b.sql"), - ("XSQBRT103", "models/b.sql"), - ("XSQBRT101", "models/c.sql"), - ), - hosts=2, - expected_plans=( - {"XSQBRT101": ["models/b.sql"], "XSQBRT102": [""], "XSQBRT103": ["models/b.sql"]}, - {"XSQBRT101": ["models/a.sql", "models/c.sql"], "XSQBRT102": [], "XSQBRT103": []}, - ), - ), - HostPlanTestCase( - description="one subject keeps the original plan", - plan={"XSQBRT101": None}, - invocations=(("XSQBRT101", "models/a.sql"),), - hosts=4, - expected_plans=({"XSQBRT101": None},), - ), - ], - ids=lambda case: case.description, -) -def test_given_planned_invocations_when_partitioning_then_split_is_deterministic( - test_case: HostPlanTestCase, -) -> None: - plans: tuple[dict[str, list[str] | None], ...] = custom_host_pool.partition_host_plan( - plan=test_case.plan, invocations=test_case.invocations, hosts=test_case.hosts - ) - - assert plans == test_case.expected_plans - - -@pytest.mark.parametrize( - "test_case", - [ - HostCancellationTestCase( - description="busy hosts stop at their next invocation", - model_count=60, - hosts=3, - timeout_millis=120_000, - delay_seconds=2.0, - expected_error_fragment="cannot inspect orders_000", - expected_max_seconds=20.0, - ), - ], - ids=lambda case: case.description, -) -def test_given_failing_host_when_others_are_busy_then_they_stop_before_later_invocations( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostCancellationTestCase -) -> None: - write_order_rules(root=tmp_path, failing=("orders_000",), delay_seconds=test_case.delay_seconds) - project: CompiledProjectView = orders_project(model_count=test_case.model_count) - launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) - monkeypatch.setattr(custom_host_pool, "_HOST_TIMEOUT_MILLIS", test_case.timeout_millis) - started: float = time.monotonic() - - with pytest.raises(RulesError) as raised: - _ = evaluate_order_rules(project=project, root=tmp_path) - - assert time.monotonic() - started < test_case.expected_max_seconds - assert test_case.expected_error_fragment in str(raised.value) - assert len(launches) == test_case.hosts + 1 - assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) - - -@pytest.mark.parametrize( - "test_case", - [ - HostCancellationTestCase( - description="busy hosts stop at their next invocation", - model_count=60, - hosts=3, - timeout_millis=120_000, - delay_seconds=2.0, - expected_error_fragment="cannot inspect orders_000", - expected_max_seconds=20.0, - ), - ], - ids=lambda case: case.description, -) -def test_given_cancelled_hosts_when_failing_then_no_process_is_signalled( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostCancellationTestCase -) -> None: - write_order_rules(root=tmp_path, failing=("orders_000",), delay_seconds=test_case.delay_seconds) - project: CompiledProjectView = orders_project(model_count=test_case.model_count) - _ = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) - signals: list[tuple[int, int]] = record_signals(monkeypatch=monkeypatch) - - with pytest.raises(RulesError) as raised: - _ = evaluate_order_rules(project=project, root=tmp_path) - - assert signals == [] - assert test_case.expected_error_fragment in str(raised.value) - - -@pytest.mark.parametrize( - "test_case", - ( - HostInputsTestCase( - description="one host waits for a slow payload", - model_count=4, - hosts=1, - write_delay_seconds=1.0, - expected_max_seconds=60.0, - ), - HostInputsTestCase( - description="split hosts wait for a slow payload", - model_count=7, - hosts=3, - write_delay_seconds=1.0, - expected_max_seconds=60.0, - ), - ), - ids=lambda case: case.description, -) -def test_given_hosts_started_before_payload_when_payload_arrives_then_findings_are_unchanged( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostInputsTestCase -) -> None: - write_order_rules(root=tmp_path) - project: CompiledProjectView = orders_project(model_count=test_case.model_count) - oracle: CustomRulesOutcome = evaluate_order_rules( - project=project, root=tmp_path, cache_enabled=False - ) - launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) - delay_inputs(monkeypatch=monkeypatch, delay_seconds=test_case.write_delay_seconds) - started: float = time.monotonic() - - delayed: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) - - assert time.monotonic() - started < test_case.expected_max_seconds - assert len(launches) == test_case.hosts - assert delayed.findings == oracle.findings - assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) - - -@pytest.mark.parametrize( - "test_case", - ( - HostInputsTestCase( - description="one waiting host is released", - model_count=4, - hosts=1, - write_delay_seconds=0.5, - expected_max_seconds=30.0, - ), - HostInputsTestCase( - description="split waiting hosts are released", - model_count=7, - hosts=3, - write_delay_seconds=0.5, - expected_max_seconds=30.0, - ), - ), - ids=lambda case: case.description, -) -def test_given_payload_write_failure_when_hosts_wait_then_original_error_is_raised_promptly( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostInputsTestCase -) -> None: - write_order_rules(root=tmp_path) - project: CompiledProjectView = orders_project(model_count=test_case.model_count) - launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) - fail_inputs(monkeypatch=monkeypatch, delay_seconds=test_case.write_delay_seconds) - started: float = time.monotonic() - - with pytest.raises(OSError, match="project payload disk is full"): - _ = evaluate_order_rules(project=project, root=tmp_path) - - assert time.monotonic() - started < test_case.expected_max_seconds - assert len(launches) == test_case.hosts - assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) - - -@pytest.mark.parametrize( - "test_case", - ( - HostInputsTestCase( - description="one host released after its input folder vanished", - model_count=4, - hosts=1, - write_delay_seconds=2.0, - expected_max_seconds=1.0, - ), - HostInputsTestCase( - description="split hosts released after their input folder vanished", - model_count=7, - hosts=3, - write_delay_seconds=2.0, - expected_max_seconds=1.0, - ), - ), - ids=lambda case: case.description, -) -def test_given_input_folder_removed_when_payload_write_fails_then_original_error_returns_promptly( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostInputsTestCase -) -> None: - write_order_rules(root=tmp_path) - project: CompiledProjectView = orders_project(model_count=test_case.model_count) - launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) - failures: list[float] = fail_inputs_without_directory( - monkeypatch=monkeypatch, delay_seconds=test_case.write_delay_seconds - ) - - with pytest.raises(OSError, match="project payload disk is full"): - _ = evaluate_order_rules(project=project, root=tmp_path) - - assert time.monotonic() - failures[0] < test_case.expected_max_seconds - assert len(launches) == test_case.hosts - assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) +from tests.unit.src.sqlbuild.rule_engine._helpers.host.helpers import touch_files @pytest.mark.parametrize( diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_host_capacity.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_host_capacity.py deleted file mode 100644 index 025e8f4757..0000000000 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_host_capacity.py +++ /dev/null @@ -1,134 +0,0 @@ -"""Host capacity honours container CPU quotas from cgroup v1 and v2.""" - -import os -from pathlib import Path - -import pytest - -from sqlbuild.rule_engine._helpers.host import custom_host_pool -from sqlbuild.rule_engine._helpers.host.host_capacity import available_cores -from tests.unit.src.sqlbuild.rule_engine._helpers.host._test_types import HostCapacityTestCase -from tests.unit.src.sqlbuild.rule_engine._helpers.host.helpers import write_cgroup_files - -_NODE_CPUS: int = 64 - - -@pytest.mark.parametrize( - "test_case", - [ - HostCapacityTestCase( - description="cgroup v2 quota on the process cgroup", - proc_cgroup="0::/kubepods/orders\n", - files={"kubepods/orders/cpu.max": "400000 100000\n"}, - expected_cores=4, - ), - HostCapacityTestCase( - description="cgroup v2 fractional quota rounds up", - proc_cgroup="0::/kubepods/orders\n", - files={"kubepods/orders/cpu.max": "150000 100000\n"}, - expected_cores=2, - ), - HostCapacityTestCase( - description="cgroup v2 ancestor quota is tighter", - proc_cgroup="0::/kubepods/orders\n", - files={ - "kubepods/cpu.max": "200000 100000\n", - "kubepods/orders/cpu.max": "800000 100000\n", - }, - expected_cores=2, - ), - HostCapacityTestCase( - description="cgroup v2 without a quota", - proc_cgroup="0::/\n", - files={"cpu.max": "max 100000\n"}, - expected_cores=_NODE_CPUS, - ), - HostCapacityTestCase( - description="cgroup v1 CFS quota", - proc_cgroup="4:cpu,cpuacct:/docker/orders\n", - files={ - "cpu,cpuacct/docker/orders/cpu.cfs_quota_us": "300000\n", - "cpu,cpuacct/docker/orders/cpu.cfs_period_us": "100000\n", - }, - expected_cores=3, - ), - HostCapacityTestCase( - description="cgroup v1 unlimited CFS quota", - proc_cgroup="4:cpu,cpuacct:/\n", - files={ - "cpu,cpuacct/cpu.cfs_quota_us": "-1\n", - "cpu,cpuacct/cpu.cfs_period_us": "100000\n", - }, - expected_cores=_NODE_CPUS, - ), - HostCapacityTestCase( - description="no cgroup files", proc_cgroup="", files={}, expected_cores=_NODE_CPUS - ), - HostCapacityTestCase( - description="non-UTF-8 cgroup v2 quota is ignored", - proc_cgroup="0::/kubepods/orders\n", - files={"kubepods/orders/cpu.max": "\udcff\udcfe 100000\n"}, - expected_cores=_NODE_CPUS, - ), - HostCapacityTestCase( - description="non-UTF-8 process cgroup is ignored", - proc_cgroup="0::/kubepods/\udcff\n", - files={"kubepods/cpu.max": "200000 100000\n"}, - expected_cores=_NODE_CPUS, - ), - HostCapacityTestCase( - description="non-UTF-8 cgroup v1 quota is ignored", - proc_cgroup="4:cpu,cpuacct:/\n", - files={ - "cpu,cpuacct/cpu.cfs_quota_us": "\udcff\n", - "cpu,cpuacct/cpu.cfs_period_us": "100000\n", - }, - expected_cores=_NODE_CPUS, - ), - ], - ids=lambda case: case.description, -) -def test_given_cgroup_quota_when_counting_cores_then_quota_caps_process_cpus( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostCapacityTestCase -) -> None: - proc_cgroup: Path = write_cgroup_files( - root=tmp_path, proc_cgroup=test_case.proc_cgroup, files=test_case.files - ) - monkeypatch.setattr(os, "process_cpu_count", lambda: _NODE_CPUS, raising=False) - - cores: int = available_cores(cgroup_root=tmp_path / "cgroup", proc_cgroup=proc_cgroup) - - assert cores == test_case.expected_cores - - -@pytest.mark.parametrize( - "test_case", - [ - HostCapacityTestCase( - description="hosts are capped on a large node", - proc_cgroup="", - files={}, - expected_cores=8, - ) - ], - ids=lambda case: case.description, -) -def test_given_many_cores_when_partitioning_then_at_most_eight_hosts_start( - monkeypatch: pytest.MonkeyPatch, test_case: HostCapacityTestCase -) -> None: - monkeypatch.setattr(custom_host_pool, "available_cores", lambda: _NODE_CPUS) - invocations: tuple[tuple[str, str], ...] = tuple( - ("XSQBRT101", f"models/orders_{index:04d}.sql") for index in range(10_000) - ) - - plans: tuple[dict[str, list[str] | None], ...] = custom_host_pool.partition_host_plan( - plan={"XSQBRT101": None}, - invocations=invocations, - hosts=custom_host_pool.host_count(len(invocations)), - ) - - assert len(plans) == test_case.expected_cores - - -if __name__ == "__main__": - raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/run/helpers.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/run/helpers.py index 20cdb0c034..064cf94d71 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/run/helpers.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/run/helpers.py @@ -2,24 +2,12 @@ import pytest +import sqlbuild._native as native from sqlbuild import _native from sqlbuild._native import NativeRulesRequest -from sqlbuild.rule_engine._helpers.engine import custom_dispatch as native from sqlbuild.rule_engine.models import NativeRulesEvaluation -def record_custom_starts(*, monkeypatch: pytest.MonkeyPatch) -> list[str]: - """Record every custom-rule start instead of starting one.""" - - started: list[str] = [] - - def record_start(**_: object) -> None: - started.append("custom") - - monkeypatch.setattr(native, "start_custom_rules", record_start) - return started - - def record_reuse(*, monkeypatch: pytest.MonkeyPatch) -> list[bool]: """Record whether each natively evaluated request was answered from the response memo.""" diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/run/test_native_rows.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/run/test_native_rows.py index ce522e41d0..e18b5f3d7f 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/run/test_native_rows.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/run/test_native_rows.py @@ -4,8 +4,8 @@ import pytest +import sqlbuild._native as native from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.rule_engine._helpers.engine import custom_dispatch as native from sqlbuild.rule_engine.exceptions import RulesError from sqlbuild.rule_engine.models import RulesCacheConfig, RulesConfig, RulesResult from tests.unit.src.sqlbuild.rule_engine._helpers.engine.helpers import evaluate_contract_rule @@ -15,7 +15,6 @@ ) from tests.unit.src.sqlbuild.rule_engine._helpers.run.helpers import ( rebuild_identity, - record_custom_starts, record_reuse, ) from tests.unit.src.sqlbuild.rule_engine.main.evaluate.helpers import build_project @@ -35,7 +34,6 @@ def test_given_config_value_the_encoder_rejects_when_building_natively_then_rule tmp_path: Path, test_case: NativeRowsEncodeErrorTestCase, ) -> None: - started: list[str] = record_custom_starts(monkeypatch=monkeypatch) project: CompiledProjectView = build_project( name="orders", relative_path="models/orders.sql", @@ -44,15 +42,13 @@ def test_given_config_value_the_encoder_rejects_when_building_natively_then_rule ) with pytest.raises(RulesError, match=test_case.expected_message): - native.evaluate_native( + native.evaluate_native_rules( project=project, config=RulesConfig(cache=RulesCacheConfig(enabled=False)), project_dir=tmp_path, catalogue=(), ) - assert started == [] - @pytest.mark.parametrize( "test_case", diff --git a/tests/unit/src/sqlbuild/rule_engine/main/evaluate/test_cache_and_suppressions.py b/tests/unit/src/sqlbuild/rule_engine/main/evaluate/test_cache_and_suppressions.py index e94a7aa128..b17e2cb6fa 100644 --- a/tests/unit/src/sqlbuild/rule_engine/main/evaluate/test_cache_and_suppressions.py +++ b/tests/unit/src/sqlbuild/rule_engine/main/evaluate/test_cache_and_suppressions.py @@ -9,9 +9,9 @@ compiled_model_record, project_with, ) +from sqlbuild._native import evaluate_project_rules as evaluate from sqlbuild.compiler.compile.types import CompiledProjectView from sqlbuild.rule_engine.exceptions import RulesError -from sqlbuild.rule_engine.main._evaluate import evaluate from sqlbuild.rule_engine.models import RuleExemption, RuleIgnore, RulesConfig, RulesResult from tests.unit.src.sqlbuild.rule_engine.main.evaluate._test_types import PolicyBehaviorTestCase from tests.unit.src.sqlbuild.rule_engine.main.evaluate.helpers import build_project diff --git a/tests/unit/src/sqlbuild/rule_engine/main/evaluate/test_evaluate.py b/tests/unit/src/sqlbuild/rule_engine/main/evaluate/test_evaluate.py index 5d9d979b7d..2cec008958 100644 --- a/tests/unit/src/sqlbuild/rule_engine/main/evaluate/test_evaluate.py +++ b/tests/unit/src/sqlbuild/rule_engine/main/evaluate/test_evaluate.py @@ -5,8 +5,8 @@ import pytest +from sqlbuild._native import evaluate_project_rules as evaluate from sqlbuild.compiler.compile.models import CompileSqlReference -from sqlbuild.rule_engine.main._evaluate import evaluate from sqlbuild.rule_engine.models import ( Finding, RulesConfig, diff --git a/tests/unit/src/sqlbuild/rule_engine/main/evaluate_rule/test_evaluate_rule.py b/tests/unit/src/sqlbuild/rule_engine/main/evaluate_rule/test_evaluate_rule.py index c728249c29..ebdf10c00f 100644 --- a/tests/unit/src/sqlbuild/rule_engine/main/evaluate_rule/test_evaluate_rule.py +++ b/tests/unit/src/sqlbuild/rule_engine/main/evaluate_rule/test_evaluate_rule.py @@ -6,10 +6,10 @@ import pytest +from sqlbuild._native import evaluate_project_rules as evaluate from sqlbuild.compiler.compile.types import CompiledProjectView from sqlbuild.rule_engine.constants import MIN_CUSTOM_RULE_TEST_CASES from sqlbuild.rule_engine.exceptions import RulesError -from sqlbuild.rule_engine.main._evaluate import evaluate from sqlbuild.rule_engine.models import RulesCacheConfig, RulesConfig, RulesResult from sqlbuild.rule_engine.types import RuleOptionValue from sqlbuild.rules.testing import RuleCase, RuleResult, evaluate_rule From e91a37e5a79bda348fb8357917d987e8d37fd4d1 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:29:22 +0100 Subject: [PATCH 04/12] test: prove custom rules split across hosts match one host --- .../sqlbuild/cli/commands/main/test_rules.py | 67 +++++++++++++++++++ 1 file changed, 67 insertions(+) diff --git a/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py b/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py index 007cdcfe63..47eb09e291 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import os from collections.abc import Callable from pathlib import Path from threading import Event @@ -20,6 +21,7 @@ from sqlbuild.compiler.sql_analysis.models import SqlLexicalSyntax from sqlbuild.lint.models import LintRunResult from tests.integration.src.sqlbuild.cli.commands.main._test_types import ( + CustomHostSplitIntegrationTestCase, CustomRuleCacheSequenceTestCase, DynamicPivotRulesIntegrationTestCase, EarlyLintInputsTestCase, @@ -1108,6 +1110,71 @@ def final_directory(*, model: Model, ctx: RuleContext) -> list[Finding]: } <= cold +@pytest.mark.parametrize( + "test_case", + [CustomHostSplitIntegrationTestCase("custom rules split across hosts", 128, 2, "XSQBRARCH001")], + ids=lambda case: case.description, +) +def test_given_custom_rules_split_across_hosts_when_compiling_then_output_matches_one_host( + test_case: CustomHostSplitIntegrationTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + (tmp_path / "sqlbuild_project.toml").write_text( + 'name = "orders"\nadapter = "duckdb"\n\n' + f'[rules]\nselect = ["{test_case.expected_code}"]\n\n' + "[rules.thresholds]\nmin_custom_rule_test_cases = 0\n", + encoding="utf-8", + ) + for index in range(test_case.model_count): + folder: str = ("final", "staging", "staging")[index % 3] + model: Path = tmp_path / "models" / folder / f"orders_{index:03d}.sql" + model.parent.mkdir(parents=True, exist_ok=True) + model.write_text( + f'MODEL (description "Orders {index}");\nSELECT {index} AS order_id\n', + encoding="utf-8", + ) + rule_file: Path = tmp_path / "rules" / "architecture.py" + rule_file.parent.mkdir() + rule_file.write_text( + """from sqlbuild.rules import Finding, Model, RuleContext, rule + +@rule( + code="XSQBRARCH001", + message="Final models must use the final directory", + remediation="Move this model beneath models/final/.", +) +def final_directory(*, model: Model, ctx: RuleContext) -> list[Finding]: + return [] if "final" in model.path.parts else [ctx.finding(subject=model)] +""", + encoding="utf-8", + ) + + def compile_payload(*, cores: int) -> dict[str, Any]: + monkeypatch.setattr(os, "process_cpu_count", lambda: cores, raising=False) + exit_code: int = main(["--project-dir", str(tmp_path), "compile", "--json", "--no-cache"]) + assert exit_code == 1 + return json.loads(capsys.readouterr().out) + + single: dict[str, Any] = compile_payload(cores=1) + rules_bulk_cache_path(project_dir=tmp_path, file_name="custom-rules.json").parent.parent.rename( + tmp_path / "single-rules-cache" + ) + split: dict[str, Any] = compile_payload(cores=test_case.hosts) + warm: dict[str, Any] = compile_payload(cores=test_case.hosts) + + assert single["compile_timings"]["rule_cache_misses"] >= test_case.model_count + assert warm["compile_timings"]["rule_cache_misses"] == 0 + assert split["diagnostics"] == single["diagnostics"] + assert len(split["diagnostics"]) == test_case.model_count - (test_case.model_count + 2) // 3 + assert warm["diagnostics"] == single["diagnostics"] + assert warm["compile_timings"]["rule_cache_hits"] >= test_case.model_count + assert (tmp_path / "single-rules-cache" / "bulk" / "custom-rules.json").read_bytes() == ( + rules_bulk_cache_path(project_dir=tmp_path, file_name="custom-rules.json").read_bytes() + ) + + @pytest.mark.parametrize( "test_case", [RulesIntegrationTestCase("custom rule is enforced and cached", 1, "XSQBRARCH001")], From 4730963eefbe2067976b80bfd4252144dd1e56ec Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:40:33 +0100 Subject: [PATCH 05/12] fix: stop split custom rule hosts on interrupt and restore multi-host coverage --- .../_helpers/rules/custom_evaluation.rs | 22 +- .../_helpers/rules/custom_host_observation.rs | 113 ++++++ .../bindings/_helpers/rules/custom_hosts.rs | 157 ++++++-- .../src/bindings/_helpers/rules/mod.rs | 1 + .../bindings/_helpers/rules/registration.rs | 2 + .../bindings/_helpers/rules/rule_dispatch.rs | 174 +++++--- .../src/custom_host_pool/main/host_count.rs | 8 +- .../tests/test_custom_host_pool.rs | 23 +- .../src/custom_host_pool/tests/test_types.rs | 1 + src/sqlbuild/_native.pyi | 14 + .../sqlbuild/cli/commands/main/_test_types.py | 22 ++ .../src/sqlbuild/cli/commands/main/helpers.py | 67 ++++ .../sqlbuild/cli/commands/main/test_rules.py | 126 +++++- .../rule_engine/host_payload/_test_types.py | 8 + .../rule_engine/host_payload/helpers.py | 133 +++++++ .../host_payload/test_host_payload_facts.py | 136 +++++++ .../rule_engine/_helpers/host/_test_types.py | 24 +- .../rule_engine/_helpers/host/helpers.py | 82 ++++ .../_helpers/host/test_custom_host_pool.py | 371 +++++++++++++++++- 19 files changed, 1358 insertions(+), 126 deletions(-) create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_observation.rs create mode 100644 tests/integration/src/sqlbuild/rule_engine/host_payload/_test_types.py create mode 100644 tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py create mode 100644 tests/integration/src/sqlbuild/rule_engine/host_payload/test_host_payload_facts.py diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_evaluation.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_evaluation.rs index 09a74008fe..ca56e429af 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_evaluation.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_evaluation.rs @@ -4,9 +4,11 @@ use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::path::PathBuf; use std::time::Instant; -use pyo3::IntoPyObjectExt; -use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyErr, PyResult, Python}; +use pyo3::prelude::{ + Bound, PyAny, PyAnyMethods, PyErr, PyModule, PyModuleMethods, PyResult, Python, +}; use pyo3::types::{PyDict, PyDictMethods, PyTuple}; +use pyo3::{IntoPyObjectExt, pyfunction, wrap_pyfunction}; use sqlbuild_cache::digest::main::sha256_hex::sha256_hex; use sqlbuild_core::json::main::dumps::dumps; use sqlbuild_core::json::models::{JsonDialect, JsonValue, OrjsonOptions}; @@ -17,7 +19,9 @@ use sqlbuild_rules::custom_rule_cache::types::ReadSet; use crate::bindings::_helpers::cache::native_store::current_code_identity; use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported}; -use crate::bindings::_helpers::rules::custom_hosts::{HostRequest, run_custom_hosts}; +use crate::bindings::_helpers::rules::custom_hosts::{ + HostCancel, HostRequest, host_payload_project, run_custom_hosts, +}; use crate::bindings::_helpers::rules::custom_rule_cache::CustomRuleCache; const NATIVE_MODULE: &str = "sqlbuild._native"; @@ -61,6 +65,7 @@ pub(crate) struct CustomRun<'a, 'py> { pub(crate) dialect: &'a str, pub(crate) verify_determinism: bool, pub(crate) implementation_fingerprints: HashMap, + pub(crate) cancel: HostCancel, } /// Reuse results whose recorded fact reads are unchanged and evaluate only the rest. @@ -151,6 +156,7 @@ pub(crate) fn evaluate_custom_rules_cached<'py>( model_paths: &model_paths, track_reads: enabled, verify_determinism: run.verify_determinism, + cancel: &run.cancel, }, &plan, )?; @@ -493,3 +499,13 @@ fn cpu_seconds(py: Python<'_>) -> PyResult { let system: f64 = times.getattr("children_system")?.extract()?; Ok(thread + user + system) } + +/// The project every custom-rule host decodes from its payload. +#[pyfunction] +fn custom_host_payload<'py>(project: &Bound<'py, PyAny>) -> PyResult> { + host_payload_project(&host_project(project)?) +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(custom_host_payload, module)?) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_observation.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_observation.rs new file mode 100644 index 0000000000..bd84c81f04 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_observation.rs @@ -0,0 +1,113 @@ +//! Test observation of custom-rule host runs: a launch log and host fan-out overrides. + +use std::path::Path; +use std::sync::{Arc, Mutex, MutexGuard}; + +use pyo3::prelude::{ + Bound, Py, PyAny, PyAnyMethods, PyModule, PyModuleMethods, PyRef, PyResult, Python, +}; +use pyo3::types::PyTuple; +use pyo3::{pyclass, pymethods}; + +use crate::bindings::_helpers::compile_driver::python_calls::imported; + +const NATIVE_MODULE: &str = "sqlbuild._native"; +/// The `sqlbuild._native` attribute tests set to a `CustomHostObserver` to observe host runs. +const OBSERVER_ATTRIBUTE: &str = "custom_host_observer"; + +/// The requests of every host launched while observed, in launch order. +pub(crate) type LaunchLog = Arc>>; + +/// Overrides host fan-out for the runs that see it and records every host they launch. +#[pyclass(frozen, module = "sqlbuild._native")] +pub(crate) struct CustomHostObserver { + cores: Option, + min_invocations_per_host: Option, + publish_inputs: Option>, + launches: LaunchLog, +} + +#[pymethods] +impl CustomHostObserver { + #[new] + #[pyo3(signature = (*, cores=None, min_invocations_per_host=None, publish_inputs=None))] + fn new( + cores: Option, + min_invocations_per_host: Option, + publish_inputs: Option>, + ) -> Self { + Self { + cores, + min_invocations_per_host, + publish_inputs, + launches: LaunchLog::default(), + } + } + + /// The request of every host launched so far, in launch order. + fn launches<'py>(&self, py: Python<'py>) -> PyResult> { + PyTuple::new(py, locked(&self.launches).clone()) + } +} + +/// What one host run takes from the observer installed when it started. +pub(crate) struct ObservedHosts { + pub(crate) cores: Option, + pub(crate) min_invocations_per_host: Option, + pub(crate) publish_inputs: Option>, + pub(crate) launches: Option, +} + +/// The installed observer's overrides and launch log; nothing when no test observes hosts. +pub(crate) fn observed_hosts(py: Python<'_>) -> PyResult { + let installed: Bound<'_, PyAny> = imported(py, NATIVE_MODULE, OBSERVER_ATTRIBUTE)?; + if installed.is_none() { + return Ok(ObservedHosts { + cores: None, + min_invocations_per_host: None, + publish_inputs: None, + launches: None, + }); + } + let observer: PyRef<'_, CustomHostObserver> = installed.extract()?; + Ok(ObservedHosts { + cores: observer.cores, + min_invocations_per_host: observer.min_invocations_per_host, + publish_inputs: observer + .publish_inputs + .as_ref() + .map(|hook| hook.clone_ref(py)), + launches: Some(Arc::clone(&observer.launches)), + }) +} + +/// Record one host launch when the run is observed. +pub(crate) fn record_launch(launches: Option<&LaunchLog>, spec: &str) { + if let Some(launches) = launches { + locked(launches).push(spec.to_owned()); + } +} + +/// Call the observer's publish hook with the payload path before the payload is written. +pub(crate) fn before_publish( + py: Python<'_>, + publish_inputs: Option<&Py>, + path: &Path, +) -> PyResult<()> { + if let Some(hook) = publish_inputs { + hook.bind(py).call1((path.to_path_buf(),))?; + } + Ok(()) +} + +fn locked(launches: &LaunchLog) -> MutexGuard<'_, Vec> { + match launches.lock() { + Ok(guard) => guard, + Err(poisoned) => poisoned.into_inner(), + } +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_class::()?; + module.add(OBSERVER_ATTRIBUTE, module.py().None()) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs index a71aa3ee79..93fe2b427e 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs @@ -2,13 +2,18 @@ use std::collections::BTreeSet; use std::path::{Path, PathBuf}; -use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::mpsc::{Receiver, Sender, channel}; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender, channel}; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use pyo3::exceptions::PyKeyboardInterrupt; use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyErr, PyResult, Python}; use pyo3::types::{PyBytes, PyBytesMethods, PyDict, PyDictMethods, PyTuple}; use serde_json::{Map, Value, json}; -use sqlbuild_rules::custom_host_pool::constants::{CGROUP_ROOT, HOST_TIMEOUT_MILLIS, PROC_CGROUP}; +use sqlbuild_rules::custom_host_pool::constants::{ + CGROUP_ROOT, HOST_TIMEOUT_MILLIS, MIN_INVOCATIONS_PER_HOST, PROC_CGROUP, +}; use sqlbuild_rules::custom_host_pool::main::available_cores::available_cores; use sqlbuild_rules::custom_host_pool::main::decode_host_response::decode_host_response; use sqlbuild_rules::custom_host_pool::main::host_count::host_count; @@ -21,6 +26,9 @@ use sqlbuild_rules::custom_host_pool::types::{HostInvocation, HostPlan}; use sqlbuild_rules::engine::main::custom_host::run_custom_host_json; use crate::bindings::_helpers::compile_driver::python_calls::imported; +use crate::bindings::_helpers::rules::custom_host_observation::{ + LaunchLog, ObservedHosts, before_publish, observed_hosts, record_launch, +}; const NATIVE_MODULE: &str = "sqlbuild._native"; const RULES_CONSTANTS_MODULE: &str = "sqlbuild.rule_engine.constants"; @@ -49,8 +57,15 @@ const PAYLOAD_MODEL_FIELDS: [&str; 18] = [ "rejected_sql_analysis_opt_out", ]; +/// How often a waiting run checks for interrupts and cancellation. +const SIGNAL_POLL: Duration = Duration::from_millis(100); +const INTERRUPTED_MESSAGE: &str = "custom rule hosts were interrupted"; + static INPUT_COUNTER: AtomicUsize = AtomicUsize::new(0); +/// Set when the run that started these hosts was interrupted; hosts stop at their next invocation. +pub(crate) type HostCancel = Arc; + /// What every host of one run is told, before its plan and cancellation marker are added. pub(crate) struct HostRequest<'a, 'py> { pub(crate) project: &'a Bound<'py, PyAny>, @@ -61,6 +76,7 @@ pub(crate) struct HostRequest<'a, 'py> { pub(crate) model_paths: &'a [String], pub(crate) track_reads: bool, pub(crate) verify_determinism: bool, + pub(crate) cancel: &'a HostCancel, } /// One host's outcome: its slice, or the `RulesError` message it failed with. @@ -76,13 +92,21 @@ pub(crate) fn run_custom_hosts( ) -> PyResult { let invocations: Vec = planned_invocations(plan, request.rules, request.model_paths); + let observed: ObservedHosts = observed_hosts(py)?; let plans: Vec = if request.verify_determinism { vec![plan.clone()] } else { + let cores: usize = match observed.cores { + Some(cores) => cores, + None => cores(py)?, + }; + let min_invocations_per_host: usize = observed + .min_invocations_per_host + .unwrap_or(MIN_INVOCATIONS_PER_HOST); partition_host_plan( plan, &invocations, - host_count(invocations.len(), cores(py)?), + host_count(invocations.len(), (cores, min_invocations_per_host)), ) }; let input_dir: PathBuf = HOST_INPUTS @@ -112,12 +136,16 @@ pub(crate) fn run_custom_hosts( json!(request.verify_determinism), ); payload.insert("track_reads".to_owned(), json!(request.track_reads)); - let launch: Launch = Launch::new(py, payload)?; + let launch: Launch = Launch::new(py, payload, observed.launches)?; + let publish = || { + before_publish(py, observed.publish_inputs.as_ref(), &input_path)?; + write_inputs(py, &input_path, request) + }; let responses: PyResult> = run_partitions( py, (&launch, plan, &plans), - &input_dir, - (|| write_inputs(py, &input_path, request), &abandoned), + (&input_dir, request.cancel), + (publish, &abandoned), ); remove_file(&input_path); remove_file(&abandoned); @@ -130,18 +158,30 @@ struct Launch { arguments: Vec, runtime_version: String, payload: Map, + launches: Option, } impl Launch { - fn new(py: Python<'_>, payload: Map) -> PyResult { + fn new( + py: Python<'_>, + payload: Map, + launches: Option, + ) -> PyResult { Ok(Self { program: imported(py, "sys", "executable")?.extract()?, arguments: vec!["-m".to_owned(), constant(py, "CUSTOM_HOST_LAUNCH_MODULE")?], runtime_version: constant(py, "CUSTOM_HOST_RUNTIME_VERSION")?, payload, + launches, }) } + /// Run one host for `spec`, recording the launch when the run is observed. + fn run(&self, spec: &str) -> HostOutcome { + record_launch(self.launches.as_ref(), spec); + host_outcome(spec) + } + /// The `run_custom_host_json` request for one host's plan and optional cancellation marker. fn spec(&self, plan: &HostPlan, cancel_marker: Option<&Path>) -> String { let mut payload: Map = self.payload.clone(); @@ -164,13 +204,17 @@ impl Launch { fn run_partitions<'py>( py: Python<'py>, (launch, plan, plans): (&Launch, &HostPlan, &[HostPlan]), - input_dir: &Path, + (input_dir, cancel): (&Path, &HostCancel), (write_inputs, abandoned): (impl FnOnce() -> PyResult<()>, &Path), ) -> PyResult> { if plans.len() == 1 { let spec: String = launch.spec(plan, None); - let outcomes: Vec = - run_hosts(py, vec![spec], (write_inputs, abandoned), None)?; + let outcomes: Vec = run_hosts( + py, + (launch, vec![spec]), + (write_inputs, abandoned), + (None, cancel), + )?; return outcomes .into_iter() .map(|outcome| outcome.map_err(|message| rules_error(py, message))) @@ -182,10 +226,15 @@ fn run_partitions<'py>( .iter() .map(|partition| launch.spec(partition, Some(&marker))) .collect(); - let outcomes: PyResult> = - run_hosts(py, specs, (write_inputs, abandoned), Some(&marker)); + let outcomes: PyResult> = run_hosts( + py, + (launch, specs), + (write_inputs, abandoned), + (Some(&marker), cancel), + ); let _removed: std::io::Result<()> = std::fs::remove_dir_all(&cancel_dir); let outcomes: Vec = outcomes?; + interrupted(py, cancel)?; if outcomes.iter().any(Result::is_err) { return single_host(py, launch, plan).map(|slice| vec![slice]); } @@ -220,16 +269,16 @@ fn run_partitions<'py>( /// Run one host for `plan` once the payload already exists. fn single_host(py: Python<'_>, launch: &Launch, plan: &HostPlan) -> PyResult { let spec: String = launch.spec(plan, None); - py.detach(|| host_outcome(&spec)) + py.detach(|| launch.run(&spec)) .map_err(|message| rules_error(py, message)) } /// Run every host on its own thread while this thread publishes the payload they wait for. fn run_hosts( py: Python<'_>, - specs: Vec, + (launch, specs): (&Launch, Vec), (write_inputs, abandoned): (impl FnOnce() -> PyResult<()>, &Path), - cancel_marker: Option<&Path>, + (cancel_marker, cancel): (Option<&Path>, &HostCancel), ) -> PyResult> { std::thread::scope(|scope| { let (sender, finished): (Sender, Receiver) = channel(); @@ -238,7 +287,7 @@ fn run_hosts( .map(|spec| { let sender: Sender = sender.clone(); scope.spawn(move || { - let outcome: HostOutcome = host_outcome(&spec); + let outcome: HostOutcome = launch.run(&spec); let status: HostStatus = outcome.as_ref().map(drop).map_err(Clone::clone); let _sent: Result<(), std::sync::mpsc::SendError> = sender.send(status); @@ -254,31 +303,69 @@ fn run_hosts( } return Err(error); } - let count: usize = handles.len(); - py.detach(move || { - for _ in 0..count { - match finished.recv() { - Ok(Ok(())) => {} - Ok(Err(_)) | Err(_) => { - if let Some(marker) = cancel_marker { - let _touched: std::io::Result<()> = touch(marker); - } - break; - } - } - } - Ok(handles + let stopped: PyResult<()> = + await_hosts(py, (finished, handles.len()), cancel_marker, cancel); + let outcomes: Vec = py.detach(move || { + handles .into_iter() .map(|handle| { handle .join() .unwrap_or_else(|_| Err("custom rule host thread panicked".to_owned())) }) - .collect()) - }) + .collect() + }); + stopped.map(|()| outcomes) }) } +/// Wait until every host finished or one failed; an interrupt cancels the rest and is returned. +fn await_hosts( + py: Python<'_>, + (finished, count): (Receiver, usize), + cancel_marker: Option<&Path>, + cancel: &HostCancel, +) -> PyResult<()> { + let finished: Mutex> = Mutex::new(finished); + let mut remaining: usize = count; + while remaining > 0 { + let received: Result = py.detach(|| match finished.lock() { + Ok(receiver) => receiver.recv_timeout(SIGNAL_POLL), + Err(_) => Err(RecvTimeoutError::Disconnected), + }); + match received { + Ok(Ok(())) => remaining -= 1, + Ok(Err(_)) | Err(RecvTimeoutError::Disconnected) => { + cancel_hosts(cancel_marker); + return Ok(()); + } + Err(RecvTimeoutError::Timeout) => { + if let Err(error) = interrupted(py, cancel) { + cancel_hosts(cancel_marker); + return Err(error); + } + } + } + } + Ok(()) +} + +/// Raise a pending interrupt, or report that the run that started these hosts was interrupted. +fn interrupted(py: Python<'_>, cancel: &HostCancel) -> PyResult<()> { + py.check_signals()?; + if cancel.load(Ordering::Acquire) { + return Err(PyKeyboardInterrupt::new_err(INTERRUPTED_MESSAGE)); + } + Ok(()) +} + +/// Stop split hosts at their next invocation; a running invocation still ends or times out. +fn cancel_hosts(cancel_marker: Option<&Path>) { + if let Some(marker) = cancel_marker { + let _touched: std::io::Result<()> = touch(marker); + } +} + /// One host run decoded, or the message `RulesError` reports for it. fn host_outcome(spec: &str) -> HostOutcome { let response: String = run_custom_host_json(spec)?; @@ -307,7 +394,9 @@ fn write_inputs(py: Python<'_>, path: &Path, request: &HostRequest<'_, '_>) -> P } /// The project hosts decode, without model state no rule fact reads. -fn host_payload_project<'py>(project: &Bound<'py, PyAny>) -> PyResult> { +pub(crate) fn host_payload_project<'py>( + project: &Bound<'py, PyAny>, +) -> PyResult> { let py: Python<'py> = project.py(); let model_class: Bound<'py, PyAny> = imported(py, COMPILE_MODELS_MODULE, "CompiledModel")?; let mut models: Vec> = Vec::new(); diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/mod.rs index 2820c2e80c..4b10fcadae 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/mod.rs @@ -2,6 +2,7 @@ pub(crate) mod builtin_catalogue; pub(crate) mod catalogue; pub(crate) mod constants; pub(crate) mod custom_evaluation; +pub(crate) mod custom_host_observation; pub(crate) mod custom_hosts; pub(crate) mod custom_rule_cache; pub(crate) mod custom_rule_evidence; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/registration.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/registration.rs index ef27bfb097..efb60e8500 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/registration.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/registration.rs @@ -64,5 +64,7 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { selection::register(module)?; results::register(module)?; guidance::register(module)?; + crate::bindings::_helpers::rules::custom_evaluation::register(module)?; + crate::bindings::_helpers::rules::custom_host_observation::register(module)?; crate::bindings::_helpers::rules::rule_dispatch::register(module) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/rule_dispatch.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/rule_dispatch.rs index cacc14bd26..04271a8ec9 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/rule_dispatch.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/rule_dispatch.rs @@ -1,10 +1,16 @@ //! Run built-in, SQL and custom rule phases on their threads and combine their findings. use std::collections::HashMap; +use std::sync::Mutex; +use std::sync::atomic::Ordering; +use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender, channel}; use std::thread::JoinHandle; +use std::time::Duration; -use pyo3::exceptions::PyRuntimeError; -use pyo3::prelude::{Bound, Py, PyAny, PyAnyMethods, PyModule, PyModuleMethods, PyResult, Python}; +use pyo3::exceptions::{PyKeyboardInterrupt, PyRuntimeError}; +use pyo3::prelude::{ + Bound, Py, PyAny, PyAnyMethods, PyErr, PyModule, PyModuleMethods, PyResult, Python, +}; use pyo3::types::{PyDict, PyDictMethods, PyFrozenSet, PyTuple}; use pyo3::{IntoPyObjectExt, pyfunction, wrap_pyfunction}; @@ -15,6 +21,7 @@ use crate::bindings::_helpers::compile_driver::python_calls::{ use crate::bindings::_helpers::rules::custom_evaluation::{ CustomRun, evaluate_custom_rules_cached, outcome, }; +use crate::bindings::_helpers::rules::custom_hosts::HostCancel; const NATIVE_MODULE: &str = "sqlbuild._native"; const RULES_MODELS_MODULE: &str = "sqlbuild.rule_engine.models"; @@ -41,6 +48,8 @@ const CUSTOM_RULES_OPTIONS: (&[&str], &[&str]) = ( &["verify_determinism", "implementation_fingerprints"], ); const RULES_THREAD: &str = "sqlbuild-rules"; +/// How often a thread waiting for a rule phase checks for interrupts. +const JOIN_POLL: Duration = Duration::from_millis(100); const EVALUATE_RULES_OPTIONS: (&[&str], &[&str]) = ( &[ "graph", @@ -69,7 +78,31 @@ const PREPARE_OPTIONS: (&[&str], &[&str]) = ( ); /// A running custom-rule evaluation, whose result is a `CustomRulesOutcome`. -type CustomHandle = JoinHandle>>; +type CustomHandle = PhaseThread; + +/// A rule phase thread and the channel that disconnects when it finishes. +struct PhaseThread { + handle: JoinHandle>>, + finished: Receiver<()>, +} + +impl PhaseThread { + fn spawn( + name: &str, + phase: impl FnOnce() -> PyResult> + Send + 'static, + ) -> PyResult { + let (sender, finished): (Sender<()>, Receiver<()>) = channel(); + let handle: JoinHandle>> = std::thread::Builder::new() + .name(name.to_owned()) + .spawn(move || { + let result: PyResult> = phase(); + drop(sender); + result + }) + .map_err(|error| PyRuntimeError::new_err(error.to_string()))?; + Ok(Self { handle, finished }) + } +} /// The rule run one native evaluation describes; every field stays bound to its Python value. struct NativeRun<'py> { @@ -80,6 +113,7 @@ struct NativeRun<'py> { dialect: String, verify_determinism: bool, initial_findings: Option>, + cancel: HostCancel, } /// How a native evaluation obtains its custom-rule outcome. @@ -141,6 +175,7 @@ fn evaluate_project_rules<'py>( dialect, verify_determinism, initial_findings: None, + cancel: HostCancel::default(), }; evaluate_native(&run, (false, None), CustomSource::Deferred) } @@ -164,6 +199,7 @@ fn evaluate_native_rules<'py>( dialect: dialect_argument(&arguments)?, verify_determinism: arguments.flag("verify_determinism")?, initial_findings: (!initial_findings.is_none()).then_some(initial_findings), + cancel: HostCancel::default(), }; evaluate_native(&run, (defer_suppressions, None), CustomSource::Deferred) } @@ -193,6 +229,7 @@ fn evaluate_custom_rules<'py>( } else { fingerprints.extract()? }, + cancel: HostCancel::default(), }) } @@ -236,6 +273,7 @@ pub(crate) fn evaluate_rules<'py>( dialect: dialect.extract()?, verify_determinism: false, initial_findings: None, + cancel: HostCancel::default(), }; let payloads: Bound<'py, PyAny> = plan.getattr("custom_payloads")?; let custom: Option = start_custom_rules(&run, &payloads)?; @@ -247,29 +285,28 @@ pub(crate) fn evaluate_rules<'py>( payloads.clone().unbind(), ]; let native_dialect: String = run.dialect.clone(); - let native_thread: JoinHandle>> = std::thread::Builder::new() - .name(RULES_THREAD.to_owned()) - .spawn(move || { - Python::attach(|py| { - let [project, config, project_dir, catalogue, payloads] = native_run; - let run: NativeRun<'_> = NativeRun { - project: project.into_bound(py), - config: config.into_bound(py), - project_dir: project_dir.into_bound(py), - catalogue: catalogue.into_bound(py), - dialect: native_dialect, - verify_determinism: false, - initial_findings: None, - }; - evaluate_native( - &run, - (true, Some(payloads.into_bound(py))), - CustomSource::Started(custom), - ) - .map(Bound::unbind) - }) + let native_cancel: HostCancel = run.cancel.clone(); + let native_thread: PhaseThread = PhaseThread::spawn(RULES_THREAD, move || { + Python::attach(|py| { + let [project, config, project_dir, catalogue, payloads] = native_run; + let run: NativeRun<'_> = NativeRun { + project: project.into_bound(py), + config: config.into_bound(py), + project_dir: project_dir.into_bound(py), + catalogue: catalogue.into_bound(py), + dialect: native_dialect, + verify_determinism: false, + initial_findings: None, + cancel: native_cancel, + }; + evaluate_native( + &run, + (true, Some(payloads.into_bound(py))), + CustomSource::Started(custom), + ) + .map(Bound::unbind) }) - .map_err(|error| PyRuntimeError::new_err(error.to_string()))?; + })?; let sql_result: PyResult> = call_with( &native(py, "run_planned_sql_rules")?, &[ @@ -281,7 +318,8 @@ pub(crate) fn evaluate_rules<'py>( ("expansion_reuse", arguments.value("expansion_reuse")?), ], ); - let result: PyResult> = joined(py, native_thread); + cancel_on_interrupt(py, &sql_result, &run.cancel); + let result: PyResult> = joined(py, native_thread, &run.cancel); let sql_result: Bound<'py, PyAny> = sql_result?; let result: Bound<'py, PyAny> = result?; call_with( @@ -398,7 +436,7 @@ fn evaluate_native<'py>( (CustomSource::Started(handle), Ok(_)) => handle, (CustomSource::Started(handle), Err(_)) => { if let Some(handle) = handle { - let _abandoned: PyResult> = joined(py, handle); + let _abandoned: PyResult> = joined(py, handle, &run.cancel); } None } @@ -408,8 +446,9 @@ fn evaluate_native<'py>( let (_, request): (Bound<'py, PyAny>, Bound<'py, PyAny>) = prepared?; let evaluated: PyResult> = native(py, "evaluate_rules_request").and_then(|evaluate| evaluate.call1((request,))); + cancel_on_interrupt(py, &evaluated, &run.cancel); let custom: PyResult> = match custom { - Some(handle) => joined(py, handle), + Some(handle) => joined(py, handle, &run.cancel), None => outcome(py, (Vec::new(), 0, 0), (0, 0)), }; let native_result: Bound<'py, PyAny> = evaluated?; @@ -579,32 +618,67 @@ fn start_custom_rules( ]; let dialect: String = run.dialect.clone(); let verify_determinism: bool = run.verify_determinism; - std::thread::Builder::new() - .name(CUSTOM_THREAD.to_owned()) - .spawn(move || { - Python::attach(|py| { - let [project, config, project_dir] = inputs; - evaluate_custom_rules_cached(CustomRun { - project: project.bind(py), - config: config.bind(py), - project_dir: project_dir.bind(py), - rules: rules.into_iter().map(|rule| rule.into_bound(py)).collect(), - dialect: &dialect, - verify_determinism, - implementation_fingerprints: fingerprints, - }) - .map(Bound::unbind) + let cancel: HostCancel = run.cancel.clone(); + PhaseThread::spawn(CUSTOM_THREAD, move || { + Python::attach(|py| { + let [project, config, project_dir] = inputs; + evaluate_custom_rules_cached(CustomRun { + project: project.bind(py), + config: config.bind(py), + project_dir: project_dir.bind(py), + rules: rules.into_iter().map(|rule| rule.into_bound(py)).collect(), + dialect: &dialect, + verify_determinism, + implementation_fingerprints: fingerprints, + cancel, }) + .map(Bound::unbind) }) - .map(Some) - .map_err(|error| PyRuntimeError::new_err(error.to_string())) + }) + .map(Some) } -/// Wait for a phase thread without holding the interpreter, then re-raise what it raised. -fn joined<'py>(py: Python<'py>, handle: CustomHandle) -> PyResult> { - py.detach(|| handle.join()) - .map_err(|_| PyRuntimeError::new_err("rules thread panicked"))? - .map(|value| value.into_bound(py)) +/// Wait for a phase thread off the interpreter; an interrupt cancels the run's hosts first. +fn joined<'py>( + py: Python<'py>, + phase: PhaseThread, + cancel: &HostCancel, +) -> PyResult> { + let finished: Mutex> = Mutex::new(phase.finished); + let mut signal: Option = None; + loop { + let waited: Result<(), RecvTimeoutError> = py.detach(|| match finished.lock() { + Ok(receiver) => receiver.recv_timeout(JOIN_POLL), + Err(_) => Err(RecvTimeoutError::Disconnected), + }); + match waited { + Ok(()) | Err(RecvTimeoutError::Disconnected) => break, + Err(RecvTimeoutError::Timeout) => { + if signal.is_none() + && let Err(error) = py.check_signals() + { + cancel.store(true, Ordering::Release); + signal = Some(error); + } + } + } + } + let joined: PyResult> = py + .detach(|| phase.handle.join()) + .map_err(|_| PyRuntimeError::new_err("rules thread panicked"))?; + if let Some(error) = signal { + return Err(error); + } + joined.map(|value| value.into_bound(py)) +} + +/// Cancel the run's custom-rule hosts when `result` is an interrupt. +fn cancel_on_interrupt(py: Python<'_>, result: &PyResult>, cancel: &HostCancel) { + if let Err(error) = result + && error.is_instance_of::(py) + { + cancel.store(true, Ordering::Release); + } } fn native<'py>(py: Python<'py>, name: &str) -> PyResult> { diff --git a/crates/sqlbuild-rules/src/custom_host_pool/main/host_count.rs b/crates/sqlbuild-rules/src/custom_host_pool/main/host_count.rs index f0eeea510a..d8da25d7de 100644 --- a/crates/sqlbuild-rules/src/custom_host_pool/main/host_count.rs +++ b/crates/sqlbuild-rules/src/custom_host_pool/main/host_count.rs @@ -1,12 +1,12 @@ //! How many hosts a run is worth. -use crate::custom_host_pool::constants::{MAX_HOSTS, MIN_INVOCATIONS_PER_HOST}; +use crate::custom_host_pool::constants::MAX_HOSTS; -/// One host per 64 invocations, bounded by usable CPUs and eight, and at least one. +/// One host per `min_invocations_per_host` invocations, bounded by cores and eight, at least one. #[must_use] -pub fn host_count(invocations: usize, cores: usize) -> usize { +pub fn host_count(invocations: usize, (cores, min_invocations_per_host): (usize, usize)) -> usize { MAX_HOSTS .min(cores) - .min(invocations / MIN_INVOCATIONS_PER_HOST) + .min(invocations / min_invocations_per_host.max(1)) .max(1) } diff --git a/crates/sqlbuild-rules/src/custom_host_pool/tests/test_custom_host_pool.rs b/crates/sqlbuild-rules/src/custom_host_pool/tests/test_custom_host_pool.rs index 84a6c451e1..f35430e473 100644 --- a/crates/sqlbuild-rules/src/custom_host_pool/tests/test_custom_host_pool.rs +++ b/crates/sqlbuild-rules/src/custom_host_pool/tests/test_custom_host_pool.rs @@ -3,6 +3,7 @@ use std::path::PathBuf; use serde_json::json; +use crate::custom_host_pool::constants::MIN_INVOCATIONS_PER_HOST; use crate::custom_host_pool::main::available_cores::available_cores; use crate::custom_host_pool::main::decode_host_response::decode_host_response; use crate::custom_host_pool::main::host_count::host_count; @@ -147,24 +148,44 @@ fn given_invocation_count_when_sizing_hosts_then_hosts_follow_work_and_cores() { description: "hosts are capped on a large node", invocations: 10_000, cores: NODE_CPUS, + min_invocations_per_host: MIN_INVOCATIONS_PER_HOST, expected_hosts: 8, }, HostCountTestCase { description: "few invocations run on one host", invocations: 63, cores: NODE_CPUS, + min_invocations_per_host: MIN_INVOCATIONS_PER_HOST, expected_hosts: 1, }, HostCountTestCase { description: "cores cap the hosts", invocations: 1_000, cores: 2, + min_invocations_per_host: MIN_INVOCATIONS_PER_HOST, expected_hosts: 2, }, + HostCountTestCase { + description: "one invocation per host splits up to the cores", + invocations: 7, + cores: 3, + min_invocations_per_host: 1, + expected_hosts: 3, + }, + HostCountTestCase { + description: "more cores than invocations runs one host per invocation", + invocations: 3, + cores: 8, + min_invocations_per_host: 1, + expected_hosts: 3, + }, ]; for test_case in test_cases { assert_eq!( - host_count(test_case.invocations, test_case.cores), + host_count( + test_case.invocations, + (test_case.cores, test_case.min_invocations_per_host) + ), test_case.expected_hosts, "{}", test_case.description diff --git a/crates/sqlbuild-rules/src/custom_host_pool/tests/test_types.rs b/crates/sqlbuild-rules/src/custom_host_pool/tests/test_types.rs index 032f6e8e11..16ea800bd3 100644 --- a/crates/sqlbuild-rules/src/custom_host_pool/tests/test_types.rs +++ b/crates/sqlbuild-rules/src/custom_host_pool/tests/test_types.rs @@ -20,6 +20,7 @@ pub(super) struct HostCountTestCase { pub(super) description: &'static str, pub(super) invocations: usize, pub(super) cores: usize, + pub(super) min_invocations_per_host: usize, pub(super) expected_hosts: usize, } diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index 0ad4c38c75..4441fe2ffb 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -6531,6 +6531,20 @@ def evaluate_custom_rules( verify_determinism: bool = False, implementation_fingerprints: dict[str, str] | None = None, ) -> CustomRulesOutcome: ... +def custom_host_payload(project: CompiledProjectView) -> CompiledProject: ... + +class CustomHostObserver: + def __init__( + self, + *, + cores: int | None = None, + min_invocations_per_host: int | None = None, + publish_inputs: Callable[[Path], object] | None = None, + ) -> None: ... + def launches(self) -> tuple[str, ...]: ... + +custom_host_observer: CustomHostObserver | None + def evaluate_rules( *, graph: ProjectGraph, diff --git a/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py b/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py index 8e0a72a9d2..755aefccff 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py @@ -174,6 +174,28 @@ class CustomHostSplitIntegrationTestCase: expected_code: str +@dataclass(frozen=True) +class CustomHostFailureIntegrationTestCase: + """A compile whose custom rule fails for some models while its hosts are split.""" + + description: str + model_count: int + hosts: int + failing_models: tuple[str, ...] + expected_failing_model: str + + +@dataclass(frozen=True) +class CustomHostInterruptIntegrationTestCase: + """A compile interrupted while split hosts run a slow custom rule.""" + + description: str + model_count: int + hosts: int + interrupt_after_seconds: float + expected_max_seconds: float + + @dataclass(frozen=True) class DollarQuotedLiteralBuildTestCase: """One model projection with dollar-quoted literals and its built rows.""" diff --git a/tests/integration/src/sqlbuild/cli/commands/main/helpers.py b/tests/integration/src/sqlbuild/cli/commands/main/helpers.py index b071a4373a..ffb3c98ed1 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/helpers.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/helpers.py @@ -1,5 +1,6 @@ from __future__ import annotations +import _thread import json import os import sys @@ -16,6 +17,7 @@ import sqlbuild._native as native_module import sqlbuild._native as plan_artifacts_module from sqlbuild._native import ( + CustomHostObserver, NativeSqlTestArtifact, PlannedStaticSqlTests, PreparedCompileArtifacts, @@ -891,3 +893,68 @@ def compile_findings( return tuple( (item["code"], item["path"], item.get("line"), item.get("column")) for item in diagnostics ) + + +CUSTOM_HOST_RULE_CODE: str = "XSQBRHOST001" + + +def write_custom_host_project(*, project_dir: Path, model_count: int, rule_body: str) -> None: + """Write numbered order models and one custom model rule running `rule_body`.""" + + (project_dir / "sqlbuild_project.toml").write_text( + 'name = "orders"\nadapter = "duckdb"\n\n' + f'[rules]\nselect = ["{CUSTOM_HOST_RULE_CODE}"]\n\n' + "[rules.thresholds]\nmin_custom_rule_test_cases = 0\n", + encoding="utf-8", + ) + for index in range(model_count): + model: Path = project_dir / "models" / f"orders_{index:03d}.sql" + model.parent.mkdir(parents=True, exist_ok=True) + model.write_text( + f'MODEL (description "Orders {index}");\nSELECT {index} AS order_id\n', + encoding="utf-8", + ) + body: str = "".join(f" {line}\n" for line in rule_body.splitlines()) + rule_file: Path = project_dir / "rules" / "orders.py" + rule_file.parent.mkdir(parents=True, exist_ok=True) + rule_file.write_text( + "from sqlbuild.rules import Finding, Model, RuleContext, rule\n\n\n" + f'@rule(code="{CUSTOM_HOST_RULE_CODE}", message="orders", remediation="Fix orders.")\n' + "def check_orders(*, model: Model, ctx: RuleContext) -> list[Finding]:\n" + f"{body}", + encoding="utf-8", + ) + + +def interrupt_main_after(delay_seconds: float) -> Callable[[Path], None]: + """Return a host publish hook that interrupts the main thread after a delay.""" + + def interrupt(path: Path) -> None: + del path + threading.Timer(delay_seconds, _thread.interrupt_main).start() + + return interrupt + + +def child_processes() -> frozenset[int]: + """Return the live child processes of this process.""" + + children: set[int] = set() + for task in (Path("/proc") / str(os.getpid()) / "task").iterdir(): + children.update(map(int, (task / "children").read_text().split())) + return frozenset(children) + + +def observe_hosts( + *, + monkeypatch: pytest.MonkeyPatch, + cores: int, + publish_inputs: Callable[[Path], None] | None = None, +) -> CustomHostObserver: + """Run custom rules on up to `cores` hosts, one invocation each, and record every launch.""" + + observer: CustomHostObserver = CustomHostObserver( + cores=cores, min_invocations_per_host=1, publish_inputs=publish_inputs + ) + monkeypatch.setattr(native_module, "custom_host_observer", observer) + return observer diff --git a/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py b/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py index 47eb09e291..3c4eceb705 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py @@ -3,7 +3,7 @@ from __future__ import annotations import json -import os +import time from collections.abc import Callable from pathlib import Path from threading import Event @@ -14,13 +14,15 @@ import sqlbuild._native as native_module import sqlbuild.compiler.compile._helpers.render.user_macros as macros_module -from sqlbuild._native import active_lint_stop, rules_bulk_cache_path +from sqlbuild._native import CustomHostObserver, active_lint_stop, rules_bulk_cache_path from sqlbuild._native import run_cli as main from sqlbuild.compiler.compile.models import DeclarationScopeBuild, LoadedMacro from sqlbuild.compiler.discovery.models import DiscoveredMacroFile, DiscoveredProjectInputs from sqlbuild.compiler.sql_analysis.models import SqlLexicalSyntax from sqlbuild.lint.models import LintRunResult from tests.integration.src.sqlbuild.cli.commands.main._test_types import ( + CustomHostFailureIntegrationTestCase, + CustomHostInterruptIntegrationTestCase, CustomHostSplitIntegrationTestCase, CustomRuleCacheSequenceTestCase, DynamicPivotRulesIntegrationTestCase, @@ -35,8 +37,12 @@ ) from tests.integration.src.sqlbuild.cli.commands.main.helpers import ( EARLY_LINT_PROJECT_FILES, + child_processes, compile_finding_keys, compile_findings, + interrupt_main_after, + observe_hosts, + write_custom_host_project, write_project_files, ) @@ -1096,13 +1102,13 @@ def final_directory(*, model: Model, ctx: RuleContext) -> list[Finding]: """, encoding="utf-8", ) + hosts: CustomHostObserver = CustomHostObserver() + monkeypatch.setattr(native_module, "custom_host_observer", hosts) cold: set[str] = compile_finding_keys(project_dir=tmp_path, capsys=capsys) warm: set[str] = compile_finding_keys(project_dir=tmp_path, capsys=capsys) - _ = main(["--project-dir", str(tmp_path), "compile", "--json"]) - timings: dict[str, int] = json.loads(capsys.readouterr().out)["compile_timings"] - assert timings["rule_cache_misses"] == 0 + assert len(hosts.launches()) == 1 assert warm == cold assert { f"models/orders.sql:{test_case.expected_code}", @@ -1112,7 +1118,7 @@ def final_directory(*, model: Model, ctx: RuleContext) -> list[Finding]: @pytest.mark.parametrize( "test_case", - [CustomHostSplitIntegrationTestCase("custom rules split across hosts", 128, 2, "XSQBRARCH001")], + [CustomHostSplitIntegrationTestCase("custom rules split across hosts", 9, 3, "XSQBRARCH001")], ids=lambda case: case.description, ) def test_given_custom_rules_split_across_hosts_when_compiling_then_output_matches_one_host( @@ -1129,7 +1135,7 @@ def test_given_custom_rules_split_across_hosts_when_compiling_then_output_matche ) for index in range(test_case.model_count): folder: str = ("final", "staging", "staging")[index % 3] - model: Path = tmp_path / "models" / folder / f"orders_{index:03d}.sql" + model: Path = tmp_path / "models" / folder / f"orders_{index:02d}.sql" model.parent.mkdir(parents=True, exist_ok=True) model.write_text( f'MODEL (description "Orders {index}");\nSELECT {index} AS order_id\n', @@ -1151,23 +1157,28 @@ def final_directory(*, model: Model, ctx: RuleContext) -> list[Finding]: encoding="utf-8", ) - def compile_payload(*, cores: int) -> dict[str, Any]: - monkeypatch.setattr(os, "process_cpu_count", lambda: cores, raising=False) + def compile_payload() -> dict[str, Any]: exit_code: int = main(["--project-dir", str(tmp_path), "compile", "--json", "--no-cache"]) assert exit_code == 1 return json.loads(capsys.readouterr().out) - single: dict[str, Any] = compile_payload(cores=1) + single_hosts: CustomHostObserver = observe_hosts(monkeypatch=monkeypatch, cores=1) + single: dict[str, Any] = compile_payload() rules_bulk_cache_path(project_dir=tmp_path, file_name="custom-rules.json").parent.parent.rename( tmp_path / "single-rules-cache" ) - split: dict[str, Any] = compile_payload(cores=test_case.hosts) - warm: dict[str, Any] = compile_payload(cores=test_case.hosts) + split_hosts: CustomHostObserver = observe_hosts(monkeypatch=monkeypatch, cores=test_case.hosts) + split: dict[str, Any] = compile_payload() + split_runs: int = len(split_hosts.launches()) + warm: dict[str, Any] = compile_payload() - assert single["compile_timings"]["rule_cache_misses"] >= test_case.model_count - assert warm["compile_timings"]["rule_cache_misses"] == 0 + assert (len(single_hosts.launches()), split_runs, len(split_hosts.launches())) == ( + 1, + test_case.hosts, + test_case.hosts, + ) assert split["diagnostics"] == single["diagnostics"] - assert len(split["diagnostics"]) == test_case.model_count - (test_case.model_count + 2) // 3 + assert len(split["diagnostics"]) == test_case.model_count - test_case.model_count // 3 assert warm["diagnostics"] == single["diagnostics"] assert warm["compile_timings"]["rule_cache_hits"] >= test_case.model_count assert (tmp_path / "single-rules-cache" / "bulk" / "custom-rules.json").read_bytes() == ( @@ -1175,6 +1186,91 @@ def compile_payload(*, cores: int) -> dict[str, Any]: ) +@pytest.mark.parametrize( + "test_case", + [ + CustomHostFailureIntegrationTestCase( + description="earliest failure lives on a later host", + model_count=9, + hosts=3, + failing_models=("orders_004", "orders_007"), + expected_failing_model="orders_004", + ), + ], + ids=lambda case: case.description, +) +def test_given_failing_custom_rule_on_split_hosts_when_compiling_then_one_host_failure_is_reported( + test_case: CustomHostFailureIntegrationTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + write_custom_host_project( + project_dir=tmp_path, + model_count=test_case.model_count, + rule_body=( + f"if model.name in {test_case.failing_models!r}:\n" + ' raise ValueError(f"cannot inspect {model.name}")\n' + "return []" + ), + ) + + def failed_compile() -> tuple[int, dict[str, Any]]: + exit_code: int = main(["--project-dir", str(tmp_path), "compile", "--json", "--no-cache"]) + return exit_code, json.loads(capsys.readouterr().out) + + single_hosts: CustomHostObserver = observe_hosts(monkeypatch=monkeypatch, cores=1) + single_exit, single = failed_compile() + split_hosts: CustomHostObserver = observe_hosts(monkeypatch=monkeypatch, cores=test_case.hosts) + split_exit, split = failed_compile() + + assert split_exit == single_exit == 1 + assert split["diagnostics"] == single["diagnostics"] + assert f"cannot inspect {test_case.expected_failing_model}" in json.dumps(split["diagnostics"]) + assert (len(single_hosts.launches()), len(split_hosts.launches())) == (1, test_case.hosts + 1) + assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) + + +@pytest.mark.parametrize( + "test_case", + [ + CustomHostInterruptIntegrationTestCase( + description="busy split hosts stop after an interrupt", + model_count=45, + hosts=3, + interrupt_after_seconds=1.5, + expected_max_seconds=12.0, + ), + ], + ids=lambda case: case.description, +) +def test_given_busy_split_hosts_when_compile_is_interrupted_then_hosts_stop_and_none_remain( + test_case: CustomHostInterruptIntegrationTestCase, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + write_custom_host_project( + project_dir=tmp_path, + model_count=test_case.model_count, + rule_body="total = 0\nfor index in range(10_000_000):\n total += index\nreturn []", + ) + children_before: frozenset[int] = child_processes() + hosts: CustomHostObserver = observe_hosts( + monkeypatch=monkeypatch, + cores=test_case.hosts, + publish_inputs=interrupt_main_after(test_case.interrupt_after_seconds), + ) + started: float = time.monotonic() + + with pytest.raises(KeyboardInterrupt): + _ = main(["--project-dir", str(tmp_path), "compile", "--json", "--no-cache"]) + + assert time.monotonic() - started < test_case.expected_max_seconds + assert len(hosts.launches()) == test_case.hosts + assert child_processes() <= children_before + assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) + + @pytest.mark.parametrize( "test_case", [RulesIntegrationTestCase("custom rule is enforced and cached", 1, "XSQBRARCH001")], diff --git a/tests/integration/src/sqlbuild/rule_engine/host_payload/_test_types.py b/tests/integration/src/sqlbuild/rule_engine/host_payload/_test_types.py new file mode 100644 index 0000000000..63224ceda7 --- /dev/null +++ b/tests/integration/src/sqlbuild/rule_engine/host_payload/_test_types.py @@ -0,0 +1,8 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class HostPayloadFactsTestCase: + description: str + project_files: dict[str, str] + expected_unpopulated_trimmed_fields: frozenset[str] diff --git a/tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py b/tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py new file mode 100644 index 0000000000..8419916287 --- /dev/null +++ b/tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py @@ -0,0 +1,133 @@ +import dataclasses +from collections.abc import Callable +from itertools import product +from pathlib import Path +from typing import Any + +import pytest + +import sqlbuild.compiler.compile.models as compile_models +from sqlbuild._native import build_compiled_project, discover_project_inputs +from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter +from sqlbuild.compiler.compile.models import CompiledModel, CompiledProject +from sqlbuild.compiler.compile.types import CompiledProjectView +from sqlbuild.rule_engine._helpers.engine import fact_replay +from sqlbuild.rule_engine._helpers.engine.fact_replay import ( + fact_outcome_digest, + shared_fact_encoder, +) +from sqlbuild.rule_engine.classes.project_tree import public_model +from sqlbuild.rule_engine.classes.rule_context import RuleFactViews, build_rule_fact_views +from sqlbuild.rule_engine.models import Model +from sqlbuild.rule_engine.types import FactKey + +_TREE_ARGUMENTS: tuple[str, ...] = ("models/marts/order_totals.sql", "models") + + +class OrdersReferenceResolver: + """External reference resolver that knows no extra relations.""" + + def validate_model_names(self, *, known_model_names: set[str]) -> None: + del known_model_names + + def extend_sql_test_model_names(self, *, known_model_names: set[str]) -> set[str]: + return known_model_names + + def extend_sql_test_source_names(self, *, known_source_names: set[str]) -> set[str]: + return known_source_names + + def extend_sql_test_seed_names(self, *, known_seed_names: set[str]) -> set[str]: + return known_seed_names + + def validate_reference( + self, + *, + ref_kind: str, + ref_name: str, + ref_package: str | None, + owner_relative_sql_path: Path, + ) -> None: + del ref_kind, ref_name, ref_package, owner_relative_sql_path + + def resolve_reference( + self, *, ref_kind: str, ref_name: str, ref_package: str | None + ) -> str | None: + del ref_kind, ref_name, ref_package + return None + + +def compile_files(*, root: Path, files: dict[str, str]) -> CompiledProjectView: + """Write and compile one DuckDB project with an external reference resolver.""" + + for relative_path, contents in files.items(): + path: Path = root / relative_path + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(contents, encoding="utf-8") + return build_compiled_project( + discovered_inputs=discover_project_inputs(project_dir=root), + adapter=DuckDbAdapter(), + external_sql_reference_resolver=OrdersReferenceResolver(), + ) + + +def dropped_fields(*, before: object, after: object) -> frozenset[str]: + """Return the fields of `after`'s dataclass whose values differ between the two snapshots.""" + + names: tuple[str, ...] = tuple(field.name for field in dataclasses.fields(type(after))) + return frozenset(filter(lambda name: getattr(before, name) != getattr(after, name), names)) + + +def every_fact_key(project: CompiledProjectView) -> tuple[FactKey, ...]: + """Return a key for every fact Rules can read, over every model of the project.""" + + models: tuple[Model, ...] = tuple(public_model(model) for model in project.models) + return ( + *((fact,) for fact in fact_replay._GLOBAL_FACTS), + *( + (fact, model.path.as_posix()) + for fact, model in product(fact_replay._MODEL_PATH_FACTS, models) + ), + *((fact, model.name) for fact, model in product(fact_replay._MODEL_NAME_FACTS, models)), + *( + (fact, *_TREE_ARGUMENTS[: fact_replay._TEXT_FACT_ARITY[fact]]) + for fact in fact_replay._TEXT_FACTS + ), + ) + + +def fact_digests(*, project: CompiledProjectView, root: Path) -> dict[FactKey, tuple[str, str]]: + """Digest every fact as hosts observe it and as whole-project digests encode it.""" + + views: RuleFactViews = build_rule_fact_views( + project=project, project_dir=root, dialect="duckdb" + ) + return { + key: ( + fact_outcome_digest(views=views, key=key), + fact_outcome_digest(views=views, key=key, shared=shared_fact_encoder()), + ) + for key in every_fact_key(project) + } + + +def record_trimmed_fields( + monkeypatch: pytest.MonkeyPatch, +) -> dict[type[object], set[str]]: + """Record replaced fields and omitted constructor fields while the host payload is trimmed.""" + + trimmed: dict[type[object], set[str]] = {CompiledProject: set(), CompiledModel: set()} + replace: Callable[..., Any] = dataclasses.replace + + def recording_replace(obj: Any, /, **changes: object) -> Any: + trimmed.setdefault(type(obj), set()).update(changes) + return replace(obj, **changes) + + def recording_model(**values: Any) -> CompiledModel: + trimmed[CompiledModel].update( + {field.name for field in dataclasses.fields(CompiledModel)} - values.keys() + ) + return CompiledModel(**values) + + monkeypatch.setattr(dataclasses, "replace", recording_replace) + monkeypatch.setattr(compile_models, "CompiledModel", recording_model) + return trimmed diff --git a/tests/integration/src/sqlbuild/rule_engine/host_payload/test_host_payload_facts.py b/tests/integration/src/sqlbuild/rule_engine/host_payload/test_host_payload_facts.py new file mode 100644 index 0000000000..5088abeffd --- /dev/null +++ b/tests/integration/src/sqlbuild/rule_engine/host_payload/test_host_payload_facts.py @@ -0,0 +1,136 @@ +"""Every Rule fact reads the same values from the trimmed host payload as from the full project.""" + +from pathlib import Path + +import pytest + +from sqlbuild._native import ( + compiled_project_record, + custom_host_payload, +) +from sqlbuild.compiler.compile.models import CompiledModel, CompiledProject +from sqlbuild.compiler.compile.types import CompiledProjectView +from sqlbuild.rule_engine.classes.project_tree import public_model +from sqlbuild.rule_engine.types import FactKey +from tests.integration.src.sqlbuild.rule_engine.host_payload._test_types import ( + HostPayloadFactsTestCase, +) +from tests.integration.src.sqlbuild.rule_engine.host_payload.helpers import ( + compile_files, + dropped_fields, + fact_digests, + record_trimmed_fields, +) + + +@pytest.mark.parametrize( + "test_case", + ( + HostPayloadFactsTestCase( + description="sources, refs, contracts, macros, declarations, tests, and audits", + project_files={ + "sqlbuild_project.toml": 'name = "orders"\nadapter = "duckdb"\n', + "sources/raw.yml": ( + "sources:\n" + " - name: raw_orders\n description: Raw orders.\n" + ' expression: "(SELECT 1 AS id, CAST(10.5 AS DOUBLE) AS amount)"\n' + " columns:\n - name: id\n type: INTEGER\n" + " - name: raw_customers\n description: Raw customers.\n" + ' expression: "(SELECT 1 AS customer_id)"\n' + ), + "constants/policy.sql": "CONSTANT (name minimum_amount, value 7);\n", + "enums/order_status.sql": "ENUM (name order_status, members [PAID, OPEN]);\n", + "models/marts/_sqlbuild/_macros/amounts.py": ( + 'def doubled(amount: str) -> str:\n return f"({amount}) * 2"\n' + ), + "models/staging/stg_orders.sql": ( + "MODEL (description 'Staged orders.',\n contract enforced,\n" + " columns (\n" + " order_id (type INTEGER, nullable false, audits [not_null]),\n" + " amount (type DOUBLE),\n status (type order_status),\n ),\n);\n\n" + "SELECT CAST(id AS INTEGER) AS order_id, amount, 'PAID' AS status\n" + 'FROM __source("raw_orders")\n' + ), + "models/marts/order_totals.sql": ( + 'MODEL (description "Order totals", materialized table);\n\n' + 'SELECT o.order_id, @doubled("o.amount") AS doubled_amount, o.status\n' + 'FROM __ref("stg_orders") AS o WHERE o.amount > 7\n' + ), + "models/marts/order_flags.sql": ( + 'MODEL (description "Order flags");\n\n' + 'SELECT o.order_id, o.missing_flag AS flag FROM __ref("stg_orders") AS o\n' + ), + "models/marts/flag_counts.sql": ( + 'MODEL (description "Flag counts");\n\n' + 'SELECT f.flag FROM __ref("order_flags") AS f\n' + ), + "hooks/sql/analyze_orders.sql": ( + 'HOOK (description "Analyze orders");\n\nSELECT 1\n' + ), + "hooks/python/order_hooks.py": ( + "from sqlbuild.hooks import hook\n\n" + "@hook\n" + "def log_orders(ctx):\n" + " '''Log orders.'''\n" + " ctx.execute_sql('INSERT INTO audit_log SELECT 1')\n" + ), + "python/loaders/order_rows.py": ( + "from sqlbuild.loaders import loader\n\n" + "@loader(columns=[{'name': 'order_id', 'type': 'INTEGER'}])\n" + "def fetch_orders(ctx):\n" + " '''Fetch orders.'''\n" + " return [{'order_id': 1}]\n" + ), + "materializations/copy_table.py": ( + "from sqlbuild.executor.custom.models import (\n" + " MaterializationContext,\n MaterializationResult,\n)\n\n\n" + "def materialize(ctx: MaterializationContext) -> MaterializationResult:\n" + ' ctx.execute_sql(f"CREATE TABLE {ctx.destination} AS {ctx.sql}")\n' + " return MaterializationResult(relation=ctx.destination)\n" + ), + "tests/scenarios/paid_orders.sql": ( + "SCENARIO (description 'Paid orders.');\n\n" + "WITH\n__source__raw_orders AS (SELECT 1 AS id, 10.5 AS amount),\n" + "__expected__stg_orders AS (\n" + " SELECT 1 AS order_id, 10.5 AS amount, 'PAID' AS status\n)\n" + "SELECT 1\n" + ), + "tests/unit/test_order_totals.sql": ( + "TEST ();\n\nWITH\n__ref__stg_orders AS (\n" + " SELECT 1 AS order_id, 10.0 AS amount\n),\n" + "__expected__order_totals AS (\n" + " SELECT 1 AS order_id, 20.0 AS doubled_amount, 'PAID' AS status\n)\n" + "SELECT 1\n" + ), + }, + expected_unpopulated_trimmed_fields=frozenset(), + ), + ), + ids=lambda case: case.description, +) +def test_given_compiled_project_when_trimming_host_payload_then_every_fact_is_unchanged( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostPayloadFactsTestCase +) -> None: + project: CompiledProjectView = compile_files(root=tmp_path, files=test_case.project_files) + trimmed_fields: dict[type[object], set[str]] = record_trimmed_fields(monkeypatch) + trimmed: CompiledProjectView = custom_host_payload(project) + monkeypatch.undo() + + full_digests: dict[FactKey, tuple[str, str]] = fact_digests(project=project, root=tmp_path) + trimmed_digests: dict[FactKey, tuple[str, str]] = fact_digests(project=trimmed, root=tmp_path) + + assert trimmed_digests == full_digests + assert tuple(map(public_model, trimmed.models)) == tuple(map(public_model, project.models)) + assert ( + dropped_fields(before=compiled_project_record(project), after=trimmed) + == trimmed_fields[CompiledProject] - test_case.expected_unpopulated_trimmed_fields + ) + assert ( + frozenset().union( + *( + dropped_fields(before=model, after=trimmed_model) + for model, trimmed_model in zip(project.models, trimmed.models, strict=True) + ) + ) + == trimmed_fields[CompiledModel] - test_case.expected_unpopulated_trimmed_fields + ) diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py index efadbee60c..c0d9afb685 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py @@ -26,17 +26,6 @@ class HostFailureTestCase: expected_failing_model: str -@dataclass(frozen=True) -class HostPlanTestCase: - """One deterministic split of planned invocations.""" - - description: str - plan: dict[str, list[str] | None] - invocations: tuple[tuple[str, str], ...] - hosts: int - expected_plans: tuple[dict[str, list[str] | None], ...] - - @dataclass(frozen=True) class HostCancellationTestCase: """A fast failure on one host while the other hosts are still busy.""" @@ -44,20 +33,21 @@ class HostCancellationTestCase: description: str model_count: int hosts: int - timeout_millis: int delay_seconds: float expected_error_fragment: str expected_max_seconds: float @dataclass(frozen=True) -class HostCapacityTestCase: - """Fake cgroup files and the CPU count they allow.""" +class HostInterruptTestCase: + """An interrupt while split hosts are still busy with slow invocations.""" description: str - proc_cgroup: str - files: dict[str, str] - expected_cores: int + model_count: int + hosts: int + delay_seconds: float + interrupt_after_seconds: float + expected_max_seconds: float @dataclass(frozen=True) diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py index 5f235cc554..827be6fac4 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py @@ -1,7 +1,10 @@ """Builders for custom-rule host partitioning tests.""" +import _thread import os import random +import shutil +import threading import time from collections.abc import Callable from operator import attrgetter @@ -11,7 +14,9 @@ import pytest +import sqlbuild._native as native_module from sqlbuild._native import ( + CustomHostObserver, build_rule_catalogue, project_with, ) @@ -109,6 +114,83 @@ def evaluate_order_rules( ) +def use_hosts( + *, + monkeypatch: pytest.MonkeyPatch, + hosts: int, + publish_inputs: Callable[[Path], object] | None = None, +) -> CustomHostObserver: + """Allow up to `hosts` host processes for any plan and record every host launch.""" + + observer: CustomHostObserver = CustomHostObserver( + cores=hosts, min_invocations_per_host=1, publish_inputs=publish_inputs + ) + monkeypatch.setattr(native_module, "custom_host_observer", observer) + return observer + + +def host_inputs_left(root: Path) -> tuple[Path, ...]: + """Return every payload, marker or folder a run left in the host-input folder.""" + + return tuple((root / "target" / "rules-cache" / "host-inputs").iterdir()) + + +def child_processes() -> frozenset[int]: + """Return the live child processes of this process.""" + + children: set[int] = set() + for task in (Path("/proc") / str(os.getpid()) / "task").iterdir(): + children.update(map(int, (task / "children").read_text().split())) + return frozenset(children) + + +def delayed_publish(delay_seconds: float) -> Callable[[Path], None]: + """Return a publish hook that lets the payload be written only after a delay.""" + + def delay(path: Path) -> None: + del path + time.sleep(delay_seconds) + + return delay + + +def failed_publish(delay_seconds: float) -> Callable[[Path], None]: + """Return a publish hook that fails to write the payload after a delay.""" + + def fail(path: Path) -> None: + del path + time.sleep(delay_seconds) + raise OSError("project payload disk is full") + + return fail + + +def failed_publish_without_directory( + delay_seconds: float, +) -> tuple[Callable[[Path], None], list[float]]: + """Return a hook that removes the input folder then fails, and when each failure happened.""" + + failures: list[float] = [] + + def fail(path: Path) -> None: + time.sleep(delay_seconds) + shutil.rmtree(path.parent) + failures.append(time.monotonic()) + raise OSError("project payload disk is full") + + return fail, failures + + +def interrupting_publish(delay_seconds: float) -> Callable[[Path], None]: + """Return a hook that interrupts the main thread once hosts have run for a while.""" + + def interrupt(path: Path) -> None: + del path + threading.Timer(delay_seconds, _thread.interrupt_main).start() + + return interrupt + + def record_signals(*, monkeypatch: pytest.MonkeyPatch) -> list[tuple[int, int]]: """Replace process signalling with a recorder so no process can be signalled.""" diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py index 54a2b919f1..1260855e7c 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py @@ -1,17 +1,384 @@ -"""The custom-rule host waits for its payload only while the parent still wants it.""" +"""Custom rules split across host processes report exactly what one host reports.""" import time from pathlib import Path import pytest +from sqlbuild._native import CustomHostObserver, rules_bulk_cache_path +from sqlbuild.compiler.compile.types import CompiledProjectView from sqlbuild.rule_engine._helpers.host import custom_host +from sqlbuild.rule_engine.constants import CUSTOM_RULES_CACHE_FILE from sqlbuild.rule_engine.exceptions import HostCancelledError, RulesError +from sqlbuild.rule_engine.models import CustomRulesOutcome from tests.unit.src.sqlbuild.rule_engine._helpers.host._test_types import ( AwaitInputsTestCase, AwaitPublishedInputsTestCase, + HostCancellationTestCase, + HostFailureTestCase, + HostInputsTestCase, + HostInterruptTestCase, + HostPartitionTestCase, ) -from tests.unit.src.sqlbuild.rule_engine._helpers.host.helpers import touch_files +from tests.unit.src.sqlbuild.rule_engine._helpers.host.helpers import ( + child_processes, + delayed_publish, + evaluate_order_rules, + failed_publish, + failed_publish_without_directory, + host_inputs_left, + interrupting_publish, + orders_project, + record_signals, + touch_files, + use_hosts, + write_order_rules, +) + + +@pytest.mark.parametrize( + "test_case", + [ + HostPartitionTestCase( + description="three hosts", model_count=7, hosts=3, expected_host_runs=3 + ), + HostPartitionTestCase( + description="more hosts than subjects", model_count=2, hosts=8, expected_host_runs=3 + ), + HostPartitionTestCase( + description="module state reruns the rule in one host", + model_count=7, + hosts=3, + expected_host_runs=4, + records_module_state=True, + ), + HostPartitionTestCase( + description="cross-model duplicate detector reruns in one host", + model_count=7, + hosts=3, + expected_host_runs=4, + detects_duplicates=True, + ), + HostPartitionTestCase( + description="uncached cross-model duplicate detector reruns in one host", + model_count=7, + hosts=3, + expected_host_runs=4, + detects_duplicates=True, + cache_enabled=False, + ), + ], + ids=lambda case: case.description, +) +def test_given_partitioned_hosts_when_evaluating_then_findings_and_cache_match_one_host( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostPartitionTestCase +) -> None: + write_order_rules( + root=tmp_path, + records_module_state=test_case.records_module_state, + detects_duplicates=test_case.detects_duplicates, + ) + project: CompiledProjectView = orders_project(model_count=test_case.model_count) + cache_path: Path = rules_bulk_cache_path( + project_dir=tmp_path, file_name=CUSTOM_RULES_CACHE_FILE + ) + cache_glob: str = cache_path.relative_to(tmp_path).as_posix() + cache_enabled: bool = test_case.cache_enabled + + single_hosts: CustomHostObserver = use_hosts(monkeypatch=monkeypatch, hosts=1) + single: CustomRulesOutcome = evaluate_order_rules( + project=project, root=tmp_path, cache_enabled=cache_enabled + ) + single_warm: CustomRulesOutcome = evaluate_order_rules( + project=project, root=tmp_path, cache_enabled=cache_enabled + ) + single_launches: int = len(single_hosts.launches()) + single_cache: bytes = b"".join(path.read_bytes() for path in tmp_path.glob(cache_glob)) + cache_path.unlink(missing_ok=True) + split_hosts: CustomHostObserver = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) + split: CustomRulesOutcome = evaluate_order_rules( + project=project, root=tmp_path, cache_enabled=cache_enabled + ) + split_cold_launches: int = len(split_hosts.launches()) + split_cache: bytes = b"".join(path.read_bytes() for path in tmp_path.glob(cache_glob)) + split_warm: CustomRulesOutcome = evaluate_order_rules( + project=project, root=tmp_path, cache_enabled=cache_enabled + ) + + assert split.findings == single.findings + assert split_warm.findings == single_warm.findings == single.findings + assert single_launches == 1 + int(not cache_enabled) + assert split_cold_launches == test_case.expected_host_runs + assert (split.cache_hits, split.cache_misses) == (single.cache_hits, single.cache_misses) + assert (split_warm.cache_hits, split_warm.cache_misses) == ( + single_warm.cache_hits, + single_warm.cache_misses, + ) + assert split_cache == single_cache + assert bool(single_cache) is cache_enabled + assert not host_inputs_left(tmp_path) + + +@pytest.mark.parametrize( + "test_case", + [ + HostPartitionTestCase( + description="edit after split", model_count=6, hosts=3, expected_host_runs=3 + ) + ], + ids=lambda case: case.description, +) +def test_given_results_cached_by_split_hosts_when_rerunning_then_only_edited_subjects_rerun( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostPartitionTestCase +) -> None: + write_order_rules(root=tmp_path) + project: CompiledProjectView = orders_project(model_count=test_case.model_count) + edited: CompiledProjectView = orders_project( + model_count=test_case.model_count, edits={"orders_003": "SELECT 2 AS order_id"} + ) + hosts: CustomHostObserver = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) + + cold: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) + cold_launches: int = len(hosts.launches()) + warm: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) + warm_launches: int = len(hosts.launches()) - cold_launches + after_edit: CustomRulesOutcome = evaluate_order_rules(project=edited, root=tmp_path) + oracle: CustomRulesOutcome = evaluate_order_rules( + project=edited, root=tmp_path, cache_enabled=False + ) + + assert cold_launches == test_case.expected_host_runs + assert warm_launches == 0 + assert (warm.cache_hits, warm.cache_misses) == (cold.cache_misses, 0) + assert warm.findings == cold.findings + assert (after_edit.cache_hits, after_edit.cache_misses) == (cold.cache_misses - 1, 1) + assert after_edit.findings == oracle.findings + + +@pytest.mark.parametrize( + "test_case", + [ + HostFailureTestCase( + description="earliest failure lives on a later host", + model_count=6, + failing_models=("orders_001", "orders_003"), + expected_failing_model="orders_001", + ), + HostFailureTestCase( + description="one failing subject", + model_count=6, + failing_models=("orders_005",), + expected_failing_model="orders_005", + ), + ], + ids=lambda case: case.description, +) +def test_given_failing_rule_on_split_hosts_when_evaluating_then_single_host_error_is_raised( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostFailureTestCase +) -> None: + write_order_rules(root=tmp_path, failing=test_case.failing_models) + project: CompiledProjectView = orders_project(model_count=test_case.model_count) + + _ = use_hosts(monkeypatch=monkeypatch, hosts=1) + with pytest.raises(RulesError) as single: + _ = evaluate_order_rules(project=project, root=tmp_path) + _ = use_hosts(monkeypatch=monkeypatch, hosts=3) + with pytest.raises(RulesError) as split: + _ = evaluate_order_rules(project=project, root=tmp_path) + + assert str(split.value) == str(single.value) + assert f"cannot inspect {test_case.expected_failing_model}" in str(split.value) + assert not host_inputs_left(tmp_path) + + +@pytest.mark.parametrize( + "test_case", + [ + HostCancellationTestCase( + description="busy hosts stop at their next invocation", + model_count=60, + hosts=3, + delay_seconds=2.0, + expected_error_fragment="cannot inspect orders_000", + expected_max_seconds=20.0, + ), + ], + ids=lambda case: case.description, +) +def test_given_failing_host_when_others_are_busy_then_they_stop_before_later_invocations( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostCancellationTestCase +) -> None: + write_order_rules(root=tmp_path, failing=("orders_000",), delay_seconds=test_case.delay_seconds) + project: CompiledProjectView = orders_project(model_count=test_case.model_count) + hosts: CustomHostObserver = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) + signals: list[tuple[int, int]] = record_signals(monkeypatch=monkeypatch) + started: float = time.monotonic() + + with pytest.raises(RulesError) as raised: + _ = evaluate_order_rules(project=project, root=tmp_path) + + assert time.monotonic() - started < test_case.expected_max_seconds + assert test_case.expected_error_fragment in str(raised.value) + assert len(hosts.launches()) == test_case.hosts + 1 + assert signals == [] + assert not host_inputs_left(tmp_path) + + +@pytest.mark.parametrize( + "test_case", + [ + HostInterruptTestCase( + description="busy split hosts stop after an interrupt", + model_count=60, + hosts=3, + delay_seconds=1.0, + interrupt_after_seconds=1.5, + expected_max_seconds=10.0, + ), + ], + ids=lambda case: case.description, +) +def test_given_busy_split_hosts_when_interrupted_then_hosts_stop_and_no_process_remains( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostInterruptTestCase +) -> None: + write_order_rules(root=tmp_path, delay_seconds=test_case.delay_seconds) + project: CompiledProjectView = orders_project(model_count=test_case.model_count) + children_before: frozenset[int] = child_processes() + hosts: CustomHostObserver = use_hosts( + monkeypatch=monkeypatch, + hosts=test_case.hosts, + publish_inputs=interrupting_publish(test_case.interrupt_after_seconds), + ) + started: float = time.monotonic() + + with pytest.raises(KeyboardInterrupt): + _ = evaluate_order_rules(project=project, root=tmp_path) + + assert time.monotonic() - started < test_case.expected_max_seconds + assert len(hosts.launches()) == test_case.hosts + assert child_processes() <= children_before + assert not host_inputs_left(tmp_path) + + +@pytest.mark.parametrize( + "test_case", + ( + HostInputsTestCase( + description="one host waits for a slow payload", + model_count=4, + hosts=1, + write_delay_seconds=1.0, + expected_max_seconds=60.0, + ), + HostInputsTestCase( + description="split hosts wait for a slow payload", + model_count=7, + hosts=3, + write_delay_seconds=1.0, + expected_max_seconds=60.0, + ), + ), + ids=lambda case: case.description, +) +def test_given_hosts_started_before_payload_when_payload_arrives_then_findings_are_unchanged( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostInputsTestCase +) -> None: + write_order_rules(root=tmp_path) + project: CompiledProjectView = orders_project(model_count=test_case.model_count) + oracle: CustomRulesOutcome = evaluate_order_rules( + project=project, root=tmp_path, cache_enabled=False + ) + hosts: CustomHostObserver = use_hosts( + monkeypatch=monkeypatch, + hosts=test_case.hosts, + publish_inputs=delayed_publish(test_case.write_delay_seconds), + ) + started: float = time.monotonic() + + delayed: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) + + assert time.monotonic() - started < test_case.expected_max_seconds + assert len(hosts.launches()) == test_case.hosts + assert delayed.findings == oracle.findings + assert not host_inputs_left(tmp_path) + + +@pytest.mark.parametrize( + "test_case", + ( + HostInputsTestCase( + description="one waiting host is released", + model_count=4, + hosts=1, + write_delay_seconds=0.5, + expected_max_seconds=30.0, + ), + HostInputsTestCase( + description="split waiting hosts are released", + model_count=7, + hosts=3, + write_delay_seconds=0.5, + expected_max_seconds=30.0, + ), + ), + ids=lambda case: case.description, +) +def test_given_payload_write_failure_when_hosts_wait_then_original_error_is_raised_promptly( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostInputsTestCase +) -> None: + write_order_rules(root=tmp_path) + project: CompiledProjectView = orders_project(model_count=test_case.model_count) + hosts: CustomHostObserver = use_hosts( + monkeypatch=monkeypatch, + hosts=test_case.hosts, + publish_inputs=failed_publish(test_case.write_delay_seconds), + ) + started: float = time.monotonic() + + with pytest.raises(OSError, match="project payload disk is full"): + _ = evaluate_order_rules(project=project, root=tmp_path) + + assert time.monotonic() - started < test_case.expected_max_seconds + assert len(hosts.launches()) == test_case.hosts + assert not host_inputs_left(tmp_path) + + +@pytest.mark.parametrize( + "test_case", + ( + HostInputsTestCase( + description="one host released after its input folder vanished", + model_count=4, + hosts=1, + write_delay_seconds=2.0, + expected_max_seconds=1.0, + ), + HostInputsTestCase( + description="split hosts released after their input folder vanished", + model_count=7, + hosts=3, + write_delay_seconds=2.0, + expected_max_seconds=1.0, + ), + ), + ids=lambda case: case.description, +) +def test_given_input_folder_removed_when_payload_write_fails_then_original_error_returns_promptly( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostInputsTestCase +) -> None: + write_order_rules(root=tmp_path) + project: CompiledProjectView = orders_project(model_count=test_case.model_count) + publish, failures = failed_publish_without_directory(test_case.write_delay_seconds) + hosts: CustomHostObserver = use_hosts( + monkeypatch=monkeypatch, hosts=test_case.hosts, publish_inputs=publish + ) + + with pytest.raises(OSError, match="project payload disk is full"): + _ = evaluate_order_rules(project=project, root=tmp_path) + + assert time.monotonic() - failures[0] < test_case.expected_max_seconds + assert len(hosts.launches()) == test_case.hosts + assert not host_inputs_left(tmp_path) @pytest.mark.parametrize( From e39c0ce73ad362fed48d0d5886c9dbe3478a3a91 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:55:27 +0100 Subject: [PATCH 06/12] refactor: validate dbt project and profile YAML natively --- .../src/dbt_profile/constants.rs | 3 + .../dbt_profile/main/dbt_profile_outputs.rs | 65 +++++++ .../dbt_profile/main/dbt_project_metadata.rs | 42 ++++ .../src/dbt_profile/main/dbt_yaml_mapping.rs | 14 ++ .../src/dbt_profile/main/mod.rs | 4 + .../main/select_dbt_profile_target.rs | 36 ++++ .../sqlbuild-config/src/dbt_profile/models.rs | 41 ++++ .../src/dbt_profile/tests/helpers.rs | 37 +++- .../src/dbt_profile/tests/test_dbt_profile.rs | 183 +++++++++++++++++- .../src/dbt_profile/tests/test_types.rs | 24 ++- .../_helpers/dbt_command/init_project.rs | 49 ++--- .../src/bindings/_helpers/dbt_interop/mod.rs | 1 + .../bindings/_helpers/dbt_interop/profile.rs | 2 +- .../dbt_interop/profile_connection.rs | 39 +--- .../_helpers/dbt_interop/profile_load.rs | 147 ++++++++++++++ python-boundary.toml | 4 +- .../integrations/dbt/_helpers/profile/load.py | 108 +---------- src/sqlbuild/integrations/dbt/models.py | 18 -- .../integrations/dbt/profile/_test_types.py | 10 + .../integrations/dbt/profile/helpers.py | 6 + .../dbt/profile/test_connection.py | 139 +++++++++++++ 21 files changed, 779 insertions(+), 193 deletions(-) create mode 100644 crates/sqlbuild-config/src/dbt_profile/main/dbt_profile_outputs.rs create mode 100644 crates/sqlbuild-config/src/dbt_profile/main/dbt_project_metadata.rs create mode 100644 crates/sqlbuild-config/src/dbt_profile/main/dbt_yaml_mapping.rs create mode 100644 crates/sqlbuild-config/src/dbt_profile/main/select_dbt_profile_target.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/profile_load.rs diff --git a/crates/sqlbuild-config/src/dbt_profile/constants.rs b/crates/sqlbuild-config/src/dbt_profile/constants.rs index 12b144478c..cbfd030d3a 100644 --- a/crates/sqlbuild-config/src/dbt_profile/constants.rs +++ b/crates/sqlbuild-config/src/dbt_profile/constants.rs @@ -48,3 +48,6 @@ pub(crate) const DATABRICKS_OAUTH_KEYS: [&str; 6] = [ "oauth_scopes", ]; pub(crate) const BIGQUERY_METHOD_ERROR_SUFFIX: &str = "is not supported by SQLBuild dbt-profile normalization yet. Supported methods: oauth, service-account"; + +/// The dbt target path used when `dbt_project.yml` sets none. +pub(crate) const DEFAULT_DBT_TARGET_PATH: &str = "target"; diff --git a/crates/sqlbuild-config/src/dbt_profile/main/dbt_profile_outputs.rs b/crates/sqlbuild-config/src/dbt_profile/main/dbt_profile_outputs.rs new file mode 100644 index 0000000000..54180cf2bf --- /dev/null +++ b/crates/sqlbuild-config/src/dbt_profile/main/dbt_profile_outputs.rs @@ -0,0 +1,65 @@ +//! The outputs and default target one `profiles.yml` profile defines. + +use sqlbuild_core::text::main::python_strip::python_strip; + +use crate::dbt_profile::errors::DbtProfileError; +use crate::dbt_profile::models::{DbtProfileOutputs, DbtYamlValue, LoadedDbtProfile}; + +/// Validate one loaded profile: a mapping with non-empty named mapping outputs and a text target. +pub fn dbt_profile_outputs( + profiles_file: &str, + profile_name: &str, + loaded: &LoadedDbtProfile, +) -> Result { + if !matches!(loaded.profile, DbtYamlValue::Mapping { .. }) { + return Err(DbtProfileError(format!( + "{profiles_file} does not define dbt profile '{profile_name}'" + ))); + } + if !matches!(loaded.outputs, DbtYamlValue::Mapping { empty: false }) { + return Err(DbtProfileError(format!( + "dbt profile '{profile_name}' must define non-empty outputs" + ))); + } + let mut outputs: Vec<(String, usize)> = Vec::new(); + for (index, (name, output)) in loaded.output_entries.iter().enumerate() { + let DbtYamlValue::Text(name) = name else { + return Err(invalid_output_name(profile_name)); + }; + let stripped: &str = python_strip(name); + if stripped.is_empty() { + return Err(invalid_output_name(profile_name)); + } + if !matches!(output, DbtYamlValue::Mapping { .. }) { + return Err(DbtProfileError(format!( + "dbt profile '{profile_name}' output '{name}' must be a mapping" + ))); + } + match outputs.iter_mut().find(|(seen, _)| seen == stripped) { + Some(seen) => seen.1 = index, + None => outputs.push((stripped.to_owned(), index)), + } + } + let default_target: Option = match &loaded.target { + DbtYamlValue::Null => None, + DbtYamlValue::Text(text) if !python_strip(text).is_empty() => { + Some(python_strip(text).to_owned()) + } + DbtYamlValue::Text(_) | DbtYamlValue::Mapping { .. } | DbtYamlValue::Other => { + return Err(DbtProfileError(format!( + "dbt profile '{profile_name}' target must be a non-empty string" + ))); + } + }; + Ok(DbtProfileOutputs { + name: profile_name.to_owned(), + default_target, + outputs, + }) +} + +fn invalid_output_name(profile_name: &str) -> DbtProfileError { + DbtProfileError(format!( + "dbt profile '{profile_name}' contains an invalid output name" + )) +} diff --git a/crates/sqlbuild-config/src/dbt_profile/main/dbt_project_metadata.rs b/crates/sqlbuild-config/src/dbt_profile/main/dbt_project_metadata.rs new file mode 100644 index 0000000000..c7479cd6a4 --- /dev/null +++ b/crates/sqlbuild-config/src/dbt_profile/main/dbt_project_metadata.rs @@ -0,0 +1,42 @@ +//! The project name, profile name and target path `dbt_project.yml` defines. + +use sqlbuild_core::text::main::python_strip::python_strip; + +use crate::dbt_profile::constants::DEFAULT_DBT_TARGET_PATH; +use crate::dbt_profile::errors::DbtProfileError; +use crate::dbt_profile::models::{DbtProjectMetadata, DbtYamlValue}; + +fn required_text(file: &str, key: &str, value: &DbtYamlValue) -> Result { + if let DbtYamlValue::Text(text) = value { + let stripped: &str = python_strip(text); + if !stripped.is_empty() { + return Ok(stripped.to_owned()); + } + } + Err(DbtProfileError(format!( + "{file} must define non-empty string '{key}'" + ))) +} + +/// Validate `name`, `profile` and `target-path` from `dbt_project.yml`, stripping each. +pub fn dbt_project_metadata( + file: &str, + (name, profile, target_path): (&DbtYamlValue, &DbtYamlValue, &DbtYamlValue), +) -> Result { + let project_name: String = required_text(file, "name", name)?; + let profile_name: String = required_text(file, "profile", profile)?; + let target_path: String = match target_path { + DbtYamlValue::Null => DEFAULT_DBT_TARGET_PATH.to_owned(), + DbtYamlValue::Text(text) if !python_strip(text).is_empty() => python_strip(text).to_owned(), + DbtYamlValue::Text(_) | DbtYamlValue::Mapping { .. } | DbtYamlValue::Other => { + return Err(DbtProfileError(format!( + "{file} target-path must be a non-empty string" + ))); + } + }; + Ok(DbtProjectMetadata { + project_name, + profile_name, + target_path, + }) +} diff --git a/crates/sqlbuild-config/src/dbt_profile/main/dbt_yaml_mapping.rs b/crates/sqlbuild-config/src/dbt_profile/main/dbt_yaml_mapping.rs new file mode 100644 index 0000000000..e382af1cd5 --- /dev/null +++ b/crates/sqlbuild-config/src/dbt_profile/main/dbt_yaml_mapping.rs @@ -0,0 +1,14 @@ +//! Whether a loaded dbt config file holds a top-level mapping. + +use crate::dbt_profile::errors::DbtProfileError; +use crate::dbt_profile::models::DbtYamlValue; + +/// Accept an empty document or a mapping; anything else is not a dbt config file. +pub fn dbt_yaml_mapping(file: &str, document: &DbtYamlValue) -> Result<(), DbtProfileError> { + match document { + DbtYamlValue::Null | DbtYamlValue::Mapping { .. } => Ok(()), + DbtYamlValue::Text(_) | DbtYamlValue::Other => Err(DbtProfileError(format!( + "{file} must contain a top-level mapping" + ))), + } +} diff --git a/crates/sqlbuild-config/src/dbt_profile/main/mod.rs b/crates/sqlbuild-config/src/dbt_profile/main/mod.rs index d483631463..640790dd39 100644 --- a/crates/sqlbuild-config/src/dbt_profile/main/mod.rs +++ b/crates/sqlbuild-config/src/dbt_profile/main/mod.rs @@ -1,2 +1,6 @@ pub mod dbt_init_project_toml; +pub mod dbt_profile_outputs; +pub mod dbt_project_metadata; +pub mod dbt_yaml_mapping; pub mod normalize_dbt_profile; +pub mod select_dbt_profile_target; diff --git a/crates/sqlbuild-config/src/dbt_profile/main/select_dbt_profile_target.rs b/crates/sqlbuild-config/src/dbt_profile/main/select_dbt_profile_target.rs new file mode 100644 index 0000000000..9ff5bee2e0 --- /dev/null +++ b/crates/sqlbuild-config/src/dbt_profile/main/select_dbt_profile_target.rs @@ -0,0 +1,36 @@ +//! The profile output a run connects with. + +use crate::dbt_profile::errors::DbtProfileError; +use crate::dbt_profile::models::DbtProfileOutputs; + +/// The requested target, else the profile default, with the entry index of its output. +pub fn select_dbt_profile_target( + profile: &DbtProfileOutputs, + requested: Option<&str>, +) -> Result<(String, usize), DbtProfileError> { + let selected: &str = match requested.filter(|target| !target.is_empty()) { + Some(target) => target, + None => profile.default_target.as_deref().ok_or_else(|| { + DbtProfileError(format!( + "dbt profile '{}' does not define a default target; pass --target", + profile.name + )) + })?, + }; + match profile.outputs.iter().find(|(name, _)| name == selected) { + Some((name, index)) => Ok((name.clone(), *index)), + None => { + let mut available: Vec<&str> = profile + .outputs + .iter() + .map(|(name, _)| name.as_str()) + .collect(); + available.sort_unstable(); + Err(DbtProfileError(format!( + "dbt profile '{}' does not define target '{selected}'. Available targets: {}", + profile.name, + available.join(", ") + ))) + } + } +} diff --git a/crates/sqlbuild-config/src/dbt_profile/models.rs b/crates/sqlbuild-config/src/dbt_profile/models.rs index 9f77770ac7..a4e5264201 100644 --- a/crates/sqlbuild-config/src/dbt_profile/models.rs +++ b/crates/sqlbuild-config/src/dbt_profile/models.rs @@ -59,3 +59,44 @@ pub struct DbtInitConfig { pub target_schema: Option, pub target_database: Option, } + +/// One value loaded from a dbt YAML file, as project and profile loading tell values apart. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum DbtYamlValue { + /// `None`, or a key the mapping does not have. + Null, + Text(String), + /// A mapping and whether it has no entries. + Mapping { + empty: bool, + }, + /// Any other value. + Other, +} + +/// The `dbt_project.yml` settings profile resolution needs. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DbtProjectMetadata { + pub project_name: String, + pub profile_name: String, + pub target_path: String, +} + +/// One profile of `profiles.yml` as loaded: the profile, its outputs, their entries and its target. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct LoadedDbtProfile { + pub profile: DbtYamlValue, + pub outputs: DbtYamlValue, + /// Each `(name, output)` of `outputs` in file order, when `outputs` is a mapping. + pub output_entries: Vec<(DbtYamlValue, DbtYamlValue)>, + pub target: DbtYamlValue, +} + +/// A validated profile: its name, default target, and each output name with its entry index. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DbtProfileOutputs { + pub name: String, + pub default_target: Option, + /// Stripped output names in first-seen order; a repeated name keeps its last entry. + pub outputs: Vec<(String, usize)>, +} diff --git a/crates/sqlbuild-config/src/dbt_profile/tests/helpers.rs b/crates/sqlbuild-config/src/dbt_profile/tests/helpers.rs index ff983129d0..9471f5bbe0 100644 --- a/crates/sqlbuild-config/src/dbt_profile/tests/helpers.rs +++ b/crates/sqlbuild-config/src/dbt_profile/tests/helpers.rs @@ -1,4 +1,6 @@ -use crate::dbt_profile::models::{DbtInitConfig, ProfileEntry, ProfileValue}; +use crate::dbt_profile::models::{ + DbtInitConfig, DbtProfileOutputs, DbtYamlValue, LoadedDbtProfile, ProfileEntry, ProfileValue, +}; pub(super) fn entry(key: &str, value: ProfileValue) -> ProfileEntry { ProfileEntry { @@ -24,3 +26,36 @@ pub(super) fn init_config(target_name: &str, target_schema: Option<&str>) -> Dbt target_database: None, } } + +pub(super) fn yaml_text(value: &str) -> DbtYamlValue { + DbtYamlValue::Text(value.to_owned()) +} + +pub(super) const YAML_MAPPING: DbtYamlValue = DbtYamlValue::Mapping { empty: false }; + +/// A loaded mapping profile whose outputs are the given `(name, output)` entries. +pub(super) fn loaded_profile( + output_entries: Vec<(DbtYamlValue, DbtYamlValue)>, + target: DbtYamlValue, +) -> LoadedDbtProfile { + LoadedDbtProfile { + profile: YAML_MAPPING, + outputs: DbtYamlValue::Mapping { + empty: output_entries.is_empty(), + }, + output_entries, + target, + } +} + +pub(super) fn profile_outputs(default_target: Option<&str>, names: &[&str]) -> DbtProfileOutputs { + DbtProfileOutputs { + name: "shop".to_owned(), + default_target: default_target.map(str::to_owned), + outputs: names + .iter() + .enumerate() + .map(|(index, name)| ((*name).to_owned(), index)) + .collect(), + } +} diff --git a/crates/sqlbuild-config/src/dbt_profile/tests/test_dbt_profile.rs b/crates/sqlbuild-config/src/dbt_profile/tests/test_dbt_profile.rs index a16119e002..cfecd8dfef 100644 --- a/crates/sqlbuild-config/src/dbt_profile/tests/test_dbt_profile.rs +++ b/crates/sqlbuild-config/src/dbt_profile/tests/test_dbt_profile.rs @@ -1,11 +1,20 @@ use crate::dbt_profile::errors::DbtProfileError; use crate::dbt_profile::main::dbt_init_project_toml::dbt_init_project_toml; +use crate::dbt_profile::main::dbt_profile_outputs::dbt_profile_outputs; +use crate::dbt_profile::main::dbt_project_metadata::dbt_project_metadata; use crate::dbt_profile::main::normalize_dbt_profile::normalize_dbt_profile; +use crate::dbt_profile::main::select_dbt_profile_target::select_dbt_profile_target; use crate::dbt_profile::models::{ - ConnectionKey, ConnectionSource, NormalizedProfile, ProfileValue, + ConnectionKey, ConnectionSource, DbtProfileOutputs, DbtProjectMetadata, DbtYamlValue, + LoadedDbtProfile, NormalizedProfile, ProfileValue, +}; +use crate::dbt_profile::tests::helpers::{ + YAML_MAPPING, entry, init_config, loaded_profile, profile_outputs, text, yaml_text, +}; +use crate::dbt_profile::tests::test_types::{ + InitTomlTestCase, NormalizeTestCase, ProfileOutputsTestCase, ProjectMetadataTestCase, + SelectTargetTestCase, }; -use crate::dbt_profile::tests::helpers::{entry, init_config, text}; -use crate::dbt_profile::tests::test_types::{InitTomlTestCase, NormalizeTestCase}; #[test] fn given_rendered_profile_outputs_when_normalizing_then_matches_python_rules() { @@ -117,3 +126,171 @@ fn given_init_values_when_writing_project_toml_then_keys_and_strings_are_escaped ); } } + +#[test] +fn given_dbt_project_values_when_reading_metadata_then_strings_are_stripped_and_required() { + let test_cases = [ + ProjectMetadataTestCase { + description: "padded values are stripped and target-path defaults", + values: [ + yaml_text(" shop "), + yaml_text("orders "), + DbtYamlValue::Null, + ], + expected_result: Ok(DbtProjectMetadata { + project_name: "shop".to_owned(), + profile_name: "orders".to_owned(), + target_path: "target".to_owned(), + }), + }, + ProjectMetadataTestCase { + description: "a blank name is rejected before the profile", + values: [yaml_text(" "), DbtYamlValue::Null, DbtYamlValue::Null], + expected_result: Err(DbtProfileError( + "dbt_project.yml must define non-empty string 'name'".to_owned(), + )), + }, + ProjectMetadataTestCase { + description: "a non-string profile is rejected", + values: [yaml_text("shop"), DbtYamlValue::Other, DbtYamlValue::Null], + expected_result: Err(DbtProfileError( + "dbt_project.yml must define non-empty string 'profile'".to_owned(), + )), + }, + ProjectMetadataTestCase { + description: "a mapping target-path is rejected", + values: [yaml_text("shop"), yaml_text("orders"), YAML_MAPPING], + expected_result: Err(DbtProfileError( + "dbt_project.yml target-path must be a non-empty string".to_owned(), + )), + }, + ]; + for test_case in test_cases { + let [name, profile, target_path] = &test_case.values; + assert_eq!( + dbt_project_metadata("dbt_project.yml", (name, profile, target_path)), + test_case.expected_result, + "{}", + test_case.description + ); + } +} + +#[test] +fn given_loaded_profile_when_validating_outputs_then_names_and_target_follow_dbt_rules() { + let test_cases = [ + ProfileOutputsTestCase { + description: "a repeated stripped name keeps its first place and last output", + loaded: loaded_profile( + vec![ + (yaml_text(" dev"), YAML_MAPPING), + (yaml_text("prod"), YAML_MAPPING), + (yaml_text("dev "), YAML_MAPPING), + ], + yaml_text(" dev "), + ), + expected_result: Ok(DbtProfileOutputs { + name: "shop".to_owned(), + default_target: Some("dev".to_owned()), + outputs: vec![("dev".to_owned(), 2), ("prod".to_owned(), 1)], + }), + }, + ProfileOutputsTestCase { + description: "a profile that is not a mapping is undefined", + loaded: LoadedDbtProfile { + profile: yaml_text("shop"), + outputs: DbtYamlValue::Null, + output_entries: Vec::new(), + target: DbtYamlValue::Null, + }, + expected_result: Err(DbtProfileError( + "profiles.yml does not define dbt profile 'shop'".to_owned(), + )), + }, + ProfileOutputsTestCase { + description: "empty outputs are rejected", + loaded: loaded_profile(Vec::new(), DbtYamlValue::Null), + expected_result: Err(DbtProfileError( + "dbt profile 'shop' must define non-empty outputs".to_owned(), + )), + }, + ProfileOutputsTestCase { + description: "a non-string output name is invalid", + loaded: loaded_profile( + vec![(DbtYamlValue::Other, YAML_MAPPING)], + DbtYamlValue::Null, + ), + expected_result: Err(DbtProfileError( + "dbt profile 'shop' contains an invalid output name".to_owned(), + )), + }, + ProfileOutputsTestCase { + description: "a scalar output names its unstripped key", + loaded: loaded_profile( + vec![(yaml_text(" dev"), yaml_text("duckdb"))], + DbtYamlValue::Null, + ), + expected_result: Err(DbtProfileError( + "dbt profile 'shop' output ' dev' must be a mapping".to_owned(), + )), + }, + ProfileOutputsTestCase { + description: "a blank target is rejected after the outputs", + loaded: loaded_profile(vec![(yaml_text("dev"), YAML_MAPPING)], yaml_text(" ")), + expected_result: Err(DbtProfileError( + "dbt profile 'shop' target must be a non-empty string".to_owned(), + )), + }, + ]; + for test_case in test_cases { + assert_eq!( + dbt_profile_outputs("profiles.yml", "shop", &test_case.loaded), + test_case.expected_result, + "{}", + test_case.description + ); + } +} + +#[test] +fn given_profile_outputs_when_selecting_target_then_requested_or_default_output_is_chosen() { + let test_cases = [ + SelectTargetTestCase { + description: "the requested target wins over the default", + profile: profile_outputs(Some("dev"), &["dev", "prod"]), + requested: Some("prod"), + expected_result: Ok(("prod".to_owned(), 1)), + }, + SelectTargetTestCase { + description: "an empty request falls back to the default", + profile: profile_outputs(Some("dev"), &["dev", "prod"]), + requested: Some(""), + expected_result: Ok(("dev".to_owned(), 0)), + }, + SelectTargetTestCase { + description: "no request and no default asks for --target", + profile: profile_outputs(None, &["dev"]), + requested: None, + expected_result: Err(DbtProfileError( + "dbt profile 'shop' does not define a default target; pass --target".to_owned(), + )), + }, + SelectTargetTestCase { + description: "an unknown target lists the sorted targets", + profile: profile_outputs(Some("qa"), &["zeta", "alpha", "Beta"]), + requested: None, + expected_result: Err(DbtProfileError( + "dbt profile 'shop' does not define target 'qa'. Available targets: Beta, alpha, zeta" + .to_owned(), + )), + }, + ]; + for test_case in test_cases { + assert_eq!( + select_dbt_profile_target(&test_case.profile, test_case.requested), + test_case.expected_result, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-config/src/dbt_profile/tests/test_types.rs b/crates/sqlbuild-config/src/dbt_profile/tests/test_types.rs index 960caa3d8a..752c57b53d 100644 --- a/crates/sqlbuild-config/src/dbt_profile/tests/test_types.rs +++ b/crates/sqlbuild-config/src/dbt_profile/tests/test_types.rs @@ -1,5 +1,8 @@ use crate::dbt_profile::errors::DbtProfileError; -use crate::dbt_profile::models::{DbtInitConfig, NormalizedProfile, ProfileEntry}; +use crate::dbt_profile::models::{ + DbtInitConfig, DbtProfileOutputs, DbtProjectMetadata, DbtYamlValue, LoadedDbtProfile, + NormalizedProfile, ProfileEntry, +}; pub(super) struct NormalizeTestCase { pub(super) description: &'static str, @@ -12,3 +15,22 @@ pub(super) struct InitTomlTestCase { pub(super) config: DbtInitConfig, pub(super) expected_toml: &'static str, } + +pub(super) struct ProjectMetadataTestCase { + pub(super) description: &'static str, + pub(super) values: [DbtYamlValue; 3], + pub(super) expected_result: Result, +} + +pub(super) struct ProfileOutputsTestCase { + pub(super) description: &'static str, + pub(super) loaded: LoadedDbtProfile, + pub(super) expected_result: Result, +} + +pub(super) struct SelectTargetTestCase { + pub(super) description: &'static str, + pub(super) profile: DbtProfileOutputs, + pub(super) requested: Option<&'static str>, + pub(super) expected_result: Result<(String, usize), DbtProfileError>, +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/init_project.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/init_project.rs index 407c1345a5..55a258af72 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/init_project.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/init_project.rs @@ -13,13 +13,16 @@ use sqlbuild_core::dbt_process::models::DbtProcessOutput; use sqlbuild_core::text::main::python_strip::python_strip; use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported}; +use crate::bindings::_helpers::dbt_interop::profile_load::{ + load_dbt_project_metadata, load_selected_dbt_profile_output, +}; use crate::bindings::_helpers::dbt_interop::runner::launch_os_error; use crate::bindings::types::CompilerDetach; +use sqlbuild_config::dbt_profile::models::DbtProjectMetadata; const NATIVE_MODULE: &str = "sqlbuild._native"; const DBT_MODELS: &str = "sqlbuild.integrations.dbt.models"; const DBT_EXCEPTIONS: &str = "sqlbuild.integrations.dbt.exceptions"; -const PROFILE_LOAD_MODULE: &str = "sqlbuild.integrations.dbt._helpers.profile.load"; const PROFILE_RENDER_MODULE: &str = "sqlbuild.integrations.dbt._helpers.profile.render"; const PROJECT_FILE: &str = "sqlbuild_project.toml"; const DEFAULT_OUTPUT_DIR: &str = "sqlbuild_project"; @@ -72,10 +75,7 @@ fn build_dbt_init_project<'py>(request: &Bound<'py, PyAny>) -> PyResult = request.getattr("progress_callbacks")?; progress(&callbacks, "start", "Inspecting dbt project and profile...")?; - let metadata: Bound<'py, PyAny> = call_with( - &imported(py, PROFILE_LOAD_MODULE, "load_dbt_project_metadata")?, - &[("project_dir", paths.dbt_project_dir.clone())], - )?; + let metadata: DbtProjectMetadata = load_dbt_project_metadata(&paths.dbt_project_dir)?; let profile_name: Bound<'py, PyAny> = profile_name(request, &metadata)?; let selected: Bound<'py, PyAny> = selected_output(&paths, request, &profile_name)?; progress(&callbacks, "complete", "Inspected dbt project and profile.")?; @@ -98,12 +98,9 @@ fn build_dbt_init_project<'py>(request: &Bound<'py, PyAny>) -> PyResult = resolved( &paths.dbt_project_dir, - &path_of(py, &metadata.getattr("target_path")?)?, + &path_of(py, &PyString::new(py, &metadata.target_path).into_any())?, )?; - let project_name: String = metadata - .getattr("project_name")? - .extract::()? - .replace('-', "_"); + let project_name: String = metadata.project_name.replace('-', "_"); let target_name: Bound<'py, PyAny> = selected.getattr("target_name")?; let toml: Bound<'py, PyAny> = call_with( &imported(py, NATIVE_MODULE, "build_dbt_init_project_toml")?, @@ -170,14 +167,7 @@ fn build_dbt_init_project<'py>(request: &Bound<'py, PyAny>) -> PyResult) -> Result<(), PyErr> { let paths: InitPaths<'_> = init_paths(request)?; - let metadata: Bound<'_, PyAny> = call_with( - &imported( - request.py(), - PROFILE_LOAD_MODULE, - "load_dbt_project_metadata", - )?, - &[("project_dir", paths.dbt_project_dir.clone())], - )?; + let metadata: DbtProjectMetadata = load_dbt_project_metadata(&paths.dbt_project_dir)?; let profile_name: Bound<'_, PyAny> = profile_name(request, &metadata)?; selected_output(&paths, request, &profile_name)?; Ok(()) @@ -197,13 +187,13 @@ fn init_paths<'py>(request: &Bound<'py, PyAny>) -> PyResult> { /// `request.profile_name or metadata.profile_name`. fn profile_name<'py>( request: &Bound<'py, PyAny>, - metadata: &Bound<'py, PyAny>, + metadata: &DbtProjectMetadata, ) -> PyResult> { let requested: Bound<'py, PyAny> = request.getattr("profile_name")?; if requested.is_truthy()? { return Ok(requested); } - metadata.getattr("profile_name") + Ok(PyString::new(request.py(), &metadata.profile_name).into_any()) } fn selected_output<'py>( @@ -211,20 +201,11 @@ fn selected_output<'py>( request: &Bound<'py, PyAny>, profile_name: &Bound<'py, PyAny>, ) -> PyResult> { - let py: Python<'py> = request.py(); - let profile: Bound<'py, PyAny> = call_with( - &imported(py, PROFILE_LOAD_MODULE, "load_raw_dbt_profile")?, - &[ - ("profiles_dir", paths.profiles_dir.clone()), - ("profile_name", profile_name.clone()), - ], - )?; - call_with( - &imported(py, PROFILE_LOAD_MODULE, "select_dbt_profile_output")?, - &[ - ("profile", profile), - ("target_name", request.getattr("target_name")?), - ], + let target_name: Option = request.getattr("target_name")?.extract()?; + load_selected_dbt_profile_output( + &paths.profiles_dir, + &profile_name.extract::()?, + target_name.as_deref(), ) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/mod.rs index 4101df8eb8..79b76ab4ca 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/mod.rs @@ -12,6 +12,7 @@ pub(crate) mod plan; pub(crate) mod plan_output; pub(crate) mod profile; pub(crate) mod profile_connection; +pub(crate) mod profile_load; pub(crate) mod reference_resolver; pub(crate) mod registration; pub(crate) mod runner; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/profile.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/profile.rs index 678ca86dcd..289cde0880 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/profile.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/profile.rs @@ -33,7 +33,7 @@ const INIT_TOML_OPTIONS: (&[&str], &[&str]) = ( &[], ); -fn raised(py: Python<'_>, error: DbtProfileError) -> PyErr { +pub(crate) fn raised(py: Python<'_>, error: DbtProfileError) -> PyErr { match imported(py, DBT_ERRORS_MODULE, "DbtProfileError") .and_then(|class| class.call1((error.0,))) { diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/profile_connection.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/profile_connection.rs index 4e451be153..1b632d3c96 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/profile_connection.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/profile_connection.rs @@ -11,9 +11,11 @@ use crate::bindings::_helpers::compile_driver::python_calls::{ KeywordArguments, call_with, imported, }; use crate::bindings::_helpers::dbt_interop::cli::resolved_path; +use crate::bindings::_helpers::dbt_interop::profile_load::{ + default_profiles_dir, load_dbt_project_metadata, load_selected_dbt_profile_output, +}; const NATIVE_MODULE: &str = "sqlbuild._native"; -const PROFILE_LOAD_MODULE: &str = "sqlbuild.integrations.dbt._helpers.profile.load"; const PROFILE_RENDER_MODULE: &str = "sqlbuild.integrations.dbt._helpers.profile.render"; const DBT_ERRORS_MODULE: &str = "sqlbuild.integrations.dbt.exceptions"; const DBT_PROFILE_SOURCE: &str = "dbt_profile"; @@ -28,7 +30,7 @@ const PROFILE_CONNECTION_OPTIONS: (&[&str], &[&str]) = ( ); const NON_EMPTY_MESSAGE: &str = "dbt_profile connection values must be non-empty strings"; -fn profile_error(py: Python<'_>, message: String) -> PyErr { +pub(crate) fn profile_error(py: Python<'_>, message: String) -> PyErr { match imported(py, DBT_ERRORS_MODULE, "DbtProfileError") .and_then(|class| class.call1((message,))) { @@ -64,10 +66,6 @@ fn or_config(value: Option, fallback: &Bound<'_, PyAny>) -> PyResult(py: Python<'py>, name: &str) -> PyResult> { - imported(py, PROFILE_LOAD_MODULE, name) -} - /// The dbt profile reference rendered and normalized into SQLBuild adapter config. fn profile_connection<'py>( sqlbuild_project_dir: &Bound<'py, PyAny>, @@ -84,31 +82,14 @@ fn profile_connection<'py>( let profiles_dir: Bound<'py, PyAny> = if profiles_dir.is_truthy()? { profiles_dir } else { - load(py, "default_profiles_dir")?.call0()? + default_profiles_dir(py)? }; - let profile_name: Bound<'py, PyAny> = match profile_name { - Some(name) => PyString::new(py, &name).into_any(), - None => call_with( - &load(py, "load_dbt_project_metadata")?, - &[("project_dir", dbt_project_dir.clone())], - )? - .getattr("profile_name")?, + let profile_name: String = match profile_name { + Some(name) => name, + None => load_dbt_project_metadata(&dbt_project_dir)?.profile_name, }; - let profile: Bound<'py, PyAny> = call_with( - &load(py, "load_raw_dbt_profile")?, - &[ - ("profiles_dir", profiles_dir.clone()), - ("profile_name", profile_name), - ], - )?; - let target_name: Bound<'py, PyAny> = target_name.map_or_else( - || py.None().into_bound(py), - |name| PyString::new(py, &name).into_any(), - ); - let selected: Bound<'py, PyAny> = call_with( - &load(py, "select_dbt_profile_output")?, - &[("profile", profile), ("target_name", target_name)], - )?; + let selected: Bound<'py, PyAny> = + load_selected_dbt_profile_output(&profiles_dir, &profile_name, target_name.as_deref())?; let resolved: Bound<'py, PyAny> = call_with( &imported( py, diff --git a/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/profile_load.rs b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/profile_load.rs new file mode 100644 index 0000000000..dbc4ebfe5e --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/profile_load.rs @@ -0,0 +1,147 @@ +//! Read the dbt project metadata and the selected profile output from the user's dbt YAML files. + +use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyResult, Python}; +use pyo3::types::{PyDict, PyDictMethods, PyString}; +use sqlbuild_config::dbt_profile::main::dbt_profile_outputs::dbt_profile_outputs; +use sqlbuild_config::dbt_profile::main::dbt_project_metadata::dbt_project_metadata; +use sqlbuild_config::dbt_profile::main::dbt_yaml_mapping::dbt_yaml_mapping; +use sqlbuild_config::dbt_profile::main::select_dbt_profile_target::select_dbt_profile_target; +use sqlbuild_config::dbt_profile::models::{ + DbtProfileOutputs, DbtProjectMetadata, DbtYamlValue, LoadedDbtProfile, +}; + +use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported}; +use crate::bindings::_helpers::dbt_interop::profile::raised; +use crate::bindings::_helpers::dbt_interop::profile_connection::profile_error; + +const PROFILE_LOAD_MODULE: &str = "sqlbuild.integrations.dbt._helpers.profile.load"; +const DBT_MODELS: &str = "sqlbuild.integrations.dbt.models"; +const DBT_PROJECT_FILE: &str = "dbt_project.yml"; +const PROFILES_FILE: &str = "profiles.yml"; +const DEFAULT_PROFILES_DIR: &str = "~/.dbt"; + +fn yaml_value(value: Option<&Bound<'_, PyAny>>) -> PyResult { + let Some(value) = value.filter(|value| !value.is_none()) else { + return Ok(DbtYamlValue::Null); + }; + if value.is_instance_of::() { + return Ok(DbtYamlValue::Text(value.extract()?)); + } + if let Some(mapping) = mapping(value.clone())? { + return Ok(DbtYamlValue::Mapping { + empty: mapping.is_empty(), + }); + } + Ok(DbtYamlValue::Other) +} + +/// `value` as a dict when it is one. +fn mapping(value: Bound<'_, PyAny>) -> PyResult>> { + if !value.is_instance_of::() { + return Ok(None); + } + Ok(Some(value.cast_into::()?)) +} + +fn get<'py>( + mapping: Option<&Bound<'py, PyDict>>, + key: &str, +) -> PyResult>> { + match mapping { + Some(mapping) => mapping.get_item(key), + None => Ok(None), + } +} + +/// The top-level mapping of one dbt config file; `None` when it holds no document. +fn yaml_mapping<'py>(file: &Bound<'py, PyAny>) -> PyResult>> { + let py: Python<'py> = file.py(); + let display: String = file.str()?.extract()?; + if !file.call_method0("exists")?.is_truthy()? { + return Err(profile_error( + py, + format!("dbt config file not found: {display}"), + )); + } + let document: Bound<'py, PyAny> = call_with( + &imported(py, PROFILE_LOAD_MODULE, "load_dbt_yaml")?, + &[("file_path", file.clone())], + )?; + dbt_yaml_mapping(&display, &yaml_value(Some(&document))?).map_err(|error| raised(py, error))?; + mapping(document) +} + +/// The project name, profile name and target path of the dbt project in `project_dir`. +pub(crate) fn load_dbt_project_metadata( + project_dir: &Bound<'_, PyAny>, +) -> PyResult { + let py: Python<'_> = project_dir.py(); + let file: Bound<'_, PyAny> = project_dir.call_method1("__truediv__", (DBT_PROJECT_FILE,))?; + let payload: Option> = yaml_mapping(&file)?; + let values: [DbtYamlValue; 3] = [ + yaml_value(get(payload.as_ref(), "name")?.as_ref())?, + yaml_value(get(payload.as_ref(), "profile")?.as_ref())?, + yaml_value(get(payload.as_ref(), "target-path")?.as_ref())?, + ]; + let [name, profile, target_path] = &values; + dbt_project_metadata( + &file.str()?.extract::()?, + (name, profile, target_path), + ) + .map_err(|error| raised(py, error)) +} + +/// The `SelectedDbtProfileOutput` of `target_name`, or the profile default, from `profiles.yml`. +pub(crate) fn load_selected_dbt_profile_output<'py>( + profiles_dir: &Bound<'py, PyAny>, + profile_name: &str, + target_name: Option<&str>, +) -> PyResult> { + let py: Python<'py> = profiles_dir.py(); + let file: Bound<'py, PyAny> = profiles_dir.call_method1("__truediv__", (PROFILES_FILE,))?; + let payload: Option> = yaml_mapping(&file)?; + let profile: Option> = get(payload.as_ref(), profile_name)?; + let profile_mapping: Option> = match &profile { + Some(value) => mapping(value.clone())?, + None => None, + }; + let outputs: Option> = get(profile_mapping.as_ref(), "outputs")?; + let output_mapping: Option> = match &outputs { + Some(value) => mapping(value.clone())?, + None => None, + }; + let mut output_values: Vec> = Vec::new(); + let mut output_entries: Vec<(DbtYamlValue, DbtYamlValue)> = Vec::new(); + if let Some(mapping) = &output_mapping { + for (name, output) in mapping.iter() { + output_entries.push((yaml_value(Some(&name))?, yaml_value(Some(&output))?)); + output_values.push(output); + } + } + let loaded: LoadedDbtProfile = LoadedDbtProfile { + profile: yaml_value(profile.as_ref())?, + outputs: yaml_value(outputs.as_ref())?, + output_entries, + target: yaml_value(get(profile_mapping.as_ref(), "target")?.as_ref())?, + }; + let validated: DbtProfileOutputs = + dbt_profile_outputs(&file.str()?.extract::()?, profile_name, &loaded) + .map_err(|error| raised(py, error))?; + let (target, index): (String, usize) = + select_dbt_profile_target(&validated, target_name).map_err(|error| raised(py, error))?; + call_with( + &imported(py, DBT_MODELS, "SelectedDbtProfileOutput")?, + &[ + ("profile_name", PyString::new(py, profile_name).into_any()), + ("target_name", PyString::new(py, &target).into_any()), + ("output", output_values[index].clone()), + ], + ) +} + +/// dbt's default profiles directory. +pub(crate) fn default_profiles_dir(py: Python<'_>) -> PyResult> { + imported(py, "pathlib", "Path")? + .call1((DEFAULT_PROFILES_DIR,))? + .call_method0("expanduser") +} diff --git a/python-boundary.toml b/python-boundary.toml index 6e3e23603d..b803e03e59 100644 --- a/python-boundary.toml +++ b/python-boundary.toml @@ -2,7 +2,7 @@ # migration work that moves to Rust. `remaining_lines` must equal the current count of that work: # a PR that moves code lowers it, and adding movable Python fails the check. # Native failures must surface as errors; `handled_native_errors` lists the reviewed exceptions. -remaining_lines = 9564 +remaining_lines = 9546 [[stays]] path = "executor" @@ -272,7 +272,7 @@ note = "Loader runtime that runs inside the executor load phase." [[stays]] path = "integrations/dbt/_helpers/profile/load.py" reason = "user-code" -note = "Loads the user's dbt_project.yml and profiles.yml with PyYAML." +note = "Rule 4: PyYAML glue only; parses dbt_project.yml and profiles.yml with yaml.safe_load and reports YAMLError as DbtProfileError. Existence, mapping, field and target validation and output selection are native (dbt_interop/profile_load.rs)." [[stays]] path = "integrations/dbt/_helpers/profile/render.py" diff --git a/src/sqlbuild/integrations/dbt/_helpers/profile/load.py b/src/sqlbuild/integrations/dbt/_helpers/profile/load.py index 87b2f851d0..3282a8dac0 100644 --- a/src/sqlbuild/integrations/dbt/_helpers/profile/load.py +++ b/src/sqlbuild/integrations/dbt/_helpers/profile/load.py @@ -1,119 +1,19 @@ -"""Load dbt project and profile YAML files.""" +"""Parse dbt project and profile YAML files with PyYAML.""" from __future__ import annotations from pathlib import Path -from typing import cast import yaml from yaml import YAMLError from sqlbuild.integrations.dbt.exceptions import DbtProfileError -from sqlbuild.integrations.dbt.models import ( - DbtProjectProfileMetadata, - RawDbtProfile, - SelectedDbtProfileOutput, -) -def load_dbt_project_metadata(*, project_dir: Path) -> DbtProjectProfileMetadata: - """Load dbt project metadata required for profile resolution.""" +def load_dbt_yaml(*, file_path: Path) -> object: + """Return the YAML document of one dbt config file.""" - project_file: Path = project_dir / "dbt_project.yml" - payload: dict[str, object] = _load_yaml_mapping(file_path=project_file) - project_name: str = _required_string(payload=payload, key="name", file_path=project_file) - profile_name: str = _required_string(payload=payload, key="profile", file_path=project_file) - target_path_value: object | None = payload.get("target-path") - target_path: str = "target" - if target_path_value is not None: - if not isinstance(target_path_value, str) or not target_path_value.strip(): - raise DbtProfileError(f"{project_file} target-path must be a non-empty string") - target_path = target_path_value.strip() - return DbtProjectProfileMetadata( - project_name=project_name, - profile_name=profile_name, - target_path=target_path, - ) - - -def load_raw_dbt_profile(*, profiles_dir: Path, profile_name: str) -> RawDbtProfile: - """Load one dbt profile from profiles.yml.""" - - profiles_file: Path = profiles_dir / "profiles.yml" - payload: dict[str, object] = _load_yaml_mapping(file_path=profiles_file) - profile_payload: object | None = payload.get(profile_name) - if not isinstance(profile_payload, dict): - raise DbtProfileError(f"{profiles_file} does not define dbt profile '{profile_name}'") - profile_mapping: dict[str, object] = cast(dict[str, object], profile_payload) - raw_outputs: object | None = profile_mapping.get("outputs") - if not isinstance(raw_outputs, dict) or not raw_outputs: - raise DbtProfileError(f"dbt profile '{profile_name}' must define non-empty outputs") - outputs: dict[str, dict[str, object]] = {} - output_name: object - output_value: object - for output_name, output_value in raw_outputs.items(): - if not isinstance(output_name, str) or not output_name.strip(): - raise DbtProfileError(f"dbt profile '{profile_name}' contains an invalid output name") - if not isinstance(output_value, dict): - raise DbtProfileError( - f"dbt profile '{profile_name}' output '{output_name}' must be a mapping" - ) - outputs[output_name.strip()] = cast(dict[str, object], output_value) - default_target_value: object | None = profile_mapping.get("target") - default_target: str | None = None - if default_target_value is not None: - if not isinstance(default_target_value, str) or not default_target_value.strip(): - raise DbtProfileError(f"dbt profile '{profile_name}' target must be a non-empty string") - default_target = default_target_value.strip() - return RawDbtProfile(name=profile_name, default_target=default_target, outputs=outputs) - - -def select_dbt_profile_output( - *, profile: RawDbtProfile, target_name: str | None -) -> SelectedDbtProfileOutput: - """Select one dbt profile output target without rendering other outputs.""" - - selected_target: str | None = target_name or profile.default_target - if selected_target is None: - raise DbtProfileError( - f"dbt profile '{profile.name}' does not define a default target; pass --target" - ) - output: dict[str, object] | None = profile.outputs.get(selected_target) - if output is None: - available: str = ", ".join(sorted(profile.outputs)) - raise DbtProfileError( - f"dbt profile '{profile.name}' does not define target '{selected_target}'. " - f"Available targets: {available}" - ) - return SelectedDbtProfileOutput( - profile_name=profile.name, - target_name=selected_target, - output=output, - ) - - -def default_profiles_dir() -> Path: - """Return dbt's default profiles directory.""" - - return Path("~/.dbt").expanduser() - - -def _load_yaml_mapping(*, file_path: Path) -> dict[str, object]: - if not file_path.exists(): - raise DbtProfileError(f"dbt config file not found: {file_path}") try: - payload: object = yaml.safe_load(file_path.read_text(encoding="utf-8")) + return yaml.safe_load(file_path.read_text(encoding="utf-8")) except YAMLError as error: raise DbtProfileError(f"{file_path} contains invalid YAML: {error}") from error - if payload is None: - return {} - if not isinstance(payload, dict): - raise DbtProfileError(f"{file_path} must contain a top-level mapping") - return cast(dict[str, object], payload) - - -def _required_string(*, payload: dict[str, object], key: str, file_path: Path) -> str: - value: object | None = payload.get(key) - if not isinstance(value, str) or not value.strip(): - raise DbtProfileError(f"{file_path} must define non-empty string '{key}'") - return value.strip() diff --git a/src/sqlbuild/integrations/dbt/models.py b/src/sqlbuild/integrations/dbt/models.py index 9c8d1aba92..8818e7c014 100644 --- a/src/sqlbuild/integrations/dbt/models.py +++ b/src/sqlbuild/integrations/dbt/models.py @@ -22,24 +22,6 @@ class DbtCliOptions: defer: bool = False -@dataclass(frozen=True) -class DbtProjectProfileMetadata: - """dbt project metadata needed for profile resolution.""" - - project_name: str - profile_name: str - target_path: str - - -@dataclass(frozen=True) -class RawDbtProfile: - """Raw dbt profile payload loaded from profiles.yml.""" - - name: str - default_target: str | None - outputs: dict[str, dict[str, object]] - - @dataclass(frozen=True) class SelectedDbtProfileOutput: """Selected dbt profile target output before rendering.""" diff --git a/tests/unit/src/sqlbuild/integrations/dbt/profile/_test_types.py b/tests/unit/src/sqlbuild/integrations/dbt/profile/_test_types.py index 61893566e4..7ea1f611b6 100644 --- a/tests/unit/src/sqlbuild/integrations/dbt/profile/_test_types.py +++ b/tests/unit/src/sqlbuild/integrations/dbt/profile/_test_types.py @@ -63,3 +63,13 @@ class DbtProfileConnectionErrorTestCase: raw_config: dict[str, object] adapter_name: str expected_error_fragment: str + + +@dataclass(frozen=True) +class DbtYamlFilesTestCase: + description: str + relative_path: str + contents: str + raw_config: dict[str, object] + expected_connection: dict[str, object] | None = None + expected_error_fragment: str | None = None diff --git a/tests/unit/src/sqlbuild/integrations/dbt/profile/helpers.py b/tests/unit/src/sqlbuild/integrations/dbt/profile/helpers.py index b7bddb52b5..644ba08b2d 100644 --- a/tests/unit/src/sqlbuild/integrations/dbt/profile/helpers.py +++ b/tests/unit/src/sqlbuild/integrations/dbt/profile/helpers.py @@ -28,3 +28,9 @@ def dbt_workspace_project_config(*, adapter_name: str) -> ProjectConfig: adapter=adapter_name, dbt=DbtConfig(project_dir="dbt_project", profiles_dir="profiles"), ) + + +def replace_dbt_workspace_file(*, root: Path, relative_path: str, contents: str) -> None: + """Replace one of the workspace's dbt YAML files.""" + + (root / relative_path).write_text(contents, encoding="utf-8") diff --git a/tests/unit/src/sqlbuild/integrations/dbt/profile/test_connection.py b/tests/unit/src/sqlbuild/integrations/dbt/profile/test_connection.py index 7bafffcf43..ba1f410299 100644 --- a/tests/unit/src/sqlbuild/integrations/dbt/profile/test_connection.py +++ b/tests/unit/src/sqlbuild/integrations/dbt/profile/test_connection.py @@ -10,12 +10,21 @@ from tests.unit.src.sqlbuild.integrations.dbt.profile._test_types import ( DbtProfileConnectionErrorTestCase, DbtProfileConnectionTestCase, + DbtYamlFilesTestCase, ) from tests.unit.src.sqlbuild.integrations.dbt.profile.helpers import ( dbt_workspace_project_config, + replace_dbt_workspace_file, write_duckdb_dbt_workspace, ) +_TWO_TARGETS: str = ( + "shop:\n target: ' dev '\n outputs:\n" + " ' dev': {type: duckdb, path: first.duckdb}\n" + " prod: {type: duckdb, path: prod.duckdb}\n" + " 'dev ': {type: duckdb, path: last.duckdb}\n" +) + @pytest.mark.parametrize( "test_case", @@ -102,3 +111,133 @@ def test_given_invalid_dbt_profile_reference_when_resolving_then_raises_profile_ ) assert captured.value.code == "C242" + + +@pytest.mark.parametrize( + "test_case", + [ + DbtYamlFilesTestCase( + description="padded and repeated output names select the last output", + relative_path="profiles/profiles.yml", + contents=_TWO_TARGETS, + raw_config={"source": "dbt_profile"}, + expected_connection={"database": "last.duckdb"}, + ), + DbtYamlFilesTestCase( + description="a requested target overrides the default", + relative_path="profiles/profiles.yml", + contents=_TWO_TARGETS, + raw_config={"source": "dbt_profile", "target": "prod"}, + expected_connection={"database": "prod.duckdb"}, + ), + DbtYamlFilesTestCase( + description="an explicit profile skips the project's profile name", + relative_path="dbt_project/dbt_project.yml", + contents="name: shop\nprofile: other\n", + raw_config={"source": "dbt_profile", "profile": "shop"}, + expected_connection={"database": "warehouse.duckdb"}, + ), + ], + ids=lambda case: case.description, +) +def test_given_dbt_yaml_files_when_resolving_dbt_profile_then_selected_output_is_used( + test_case: DbtYamlFilesTestCase, tmp_path: Path +) -> None: + write_duckdb_dbt_workspace(tmp_path) + replace_dbt_workspace_file( + root=tmp_path, relative_path=test_case.relative_path, contents=test_case.contents + ) + + resolved: NormalizedDbtProfileConnection | None = resolve_raw_dbt_profile_connection( + raw_config=test_case.raw_config, + project_dir=tmp_path, + adapter_name="duckdb", + project_config=dbt_workspace_project_config(adapter_name="duckdb"), + ) + + assert resolved is not None + assert resolved.connection == test_case.expected_connection + + +@pytest.mark.parametrize( + "test_case", + [ + DbtYamlFilesTestCase( + description="a project file that is not a mapping is rejected", + relative_path="dbt_project/dbt_project.yml", + contents="- shop\n", + raw_config={"source": "dbt_profile"}, + expected_error_fragment="dbt_project.yml must contain a top-level mapping", + ), + DbtYamlFilesTestCase( + description="invalid project YAML names the file", + relative_path="dbt_project/dbt_project.yml", + contents="name: [shop\n", + raw_config={"source": "dbt_profile"}, + expected_error_fragment="dbt_project.yml contains invalid YAML", + ), + DbtYamlFilesTestCase( + description="a project without a profile name is rejected", + relative_path="dbt_project/dbt_project.yml", + contents="name: shop\n", + raw_config={"source": "dbt_profile"}, + expected_error_fragment="must define non-empty string 'profile'", + ), + DbtYamlFilesTestCase( + description="an undefined profile is rejected", + relative_path="profiles/profiles.yml", + contents="other:\n outputs:\n dev: {type: duckdb}\n", + raw_config={"source": "dbt_profile"}, + expected_error_fragment="does not define dbt profile 'shop'", + ), + DbtYamlFilesTestCase( + description="empty outputs are rejected", + relative_path="profiles/profiles.yml", + contents="shop:\n target: dev\n outputs: {}\n", + raw_config={"source": "dbt_profile"}, + expected_error_fragment="dbt profile 'shop' must define non-empty outputs", + ), + DbtYamlFilesTestCase( + description="a non-string output name is rejected", + relative_path="profiles/profiles.yml", + contents="shop:\n target: dev\n outputs:\n 1: {type: duckdb}\n", + raw_config={"source": "dbt_profile"}, + expected_error_fragment="dbt profile 'shop' contains an invalid output name", + ), + DbtYamlFilesTestCase( + description="no default target asks for one", + relative_path="profiles/profiles.yml", + contents="shop:\n outputs:\n dev: {type: duckdb}\n", + raw_config={"source": "dbt_profile"}, + expected_error_fragment="does not define a default target; pass --target", + ), + DbtYamlFilesTestCase( + description="an unknown target lists the sorted targets", + relative_path="profiles/profiles.yml", + contents="shop:\n target: qa\n outputs:\n zeta: {type: duckdb}\n" + " alpha: {type: duckdb}\n", + raw_config={"source": "dbt_profile"}, + expected_error_fragment="Available targets: alpha, zeta", + ), + ], + ids=lambda case: case.description, +) +def test_given_malformed_dbt_yaml_files_when_resolving_dbt_profile_then_raises_profile_error( + test_case: DbtYamlFilesTestCase, tmp_path: Path +) -> None: + write_duckdb_dbt_workspace(tmp_path) + replace_dbt_workspace_file( + root=tmp_path, relative_path=test_case.relative_path, contents=test_case.contents + ) + + with pytest.raises(DbtProfileError) as captured: + resolve_raw_dbt_profile_connection( + raw_config=test_case.raw_config, + project_dir=tmp_path, + adapter_name="duckdb", + project_config=dbt_workspace_project_config(adapter_name="duckdb"), + ) + + assert test_case.expected_error_fragment is not None + assert test_case.expected_error_fragment in str(captured.value) + assert captured.value.code == "C242" From ab9ea0ef5372893c186bd44b738c4351a429e04e Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 15 Oct 2026 12:26:46 +0100 Subject: [PATCH 07/12] chore: drop stale imports after orchestration integration --- src/sqlbuild/_native.pyi | 31 +++---------------- .../main/test_compile_phase_timings.py | 2 -- .../rule_engine/_helpers/engine/helpers.py | 12 +++---- 3 files changed, 8 insertions(+), 37 deletions(-) diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index 4441fe2ffb..792d2d09dc 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -258,9 +258,7 @@ from sqlbuild.compiler.scopes.models import ( UsageRecord, VisibilityResolution, ) -from sqlbuild.compiler.scopes.types import ( - DeclarationKind, -) +from sqlbuild.compiler.scopes.types import DeclarationKind from sqlbuild.compiler.source_freshness.models import ( DirectSourceFreshnessPlanningResult, SourceFreshnessRecord, @@ -275,18 +273,9 @@ from sqlbuild.compiler.sql_test_glue.models import ( NativeSqlTestChainRequest, NativeSqlTestPlanningRequest, ) -from sqlbuild.cost.models import ( - CostResourceContext, - RunCostSummary, - StatementExecutionTelemetry, -) +from sqlbuild.cost.models import CostResourceContext, RunCostSummary, StatementExecutionTelemetry from sqlbuild.cost.types import CostStatus -from sqlbuild.cursor_algebra.models import ( - AlignedInterval, - DateValue, - Duration, - TimestampValue, -) +from sqlbuild.cursor_algebra.models import AlignedInterval, DateValue, Duration, TimestampValue from sqlbuild.cursor_algebra.types import Bound, BoundSentinel, CursorScalar from sqlbuild.diagnostics.models import ProcessResourceUsage from sqlbuild.executor.auditing.models import AuditExecutionResult @@ -312,24 +301,12 @@ from sqlbuild.integrations.dbt.models import ( DbtCommandResult, DbtInitRequest, DbtInitResult, - DbtInteropCommandArgs, - DbtInteropCompiledProject, - DbtInteropExecutionRequest, - DbtInteropInvocation, - DbtInteropParsedArgs, - DbtInteropPlan, - DbtInteropPlanResolution, - DbtInteropRoutedArgs, - DbtInteropSelectionResult, - DbtLsNode, DbtLsResult, - DbtManifestIndex, - DbtManifestModel, DbtNodeExecutionResult, NormalizedDbtProfileConnection, ResolvedDbtProfileOutput, ) -from sqlbuild.integrations.dbt.types import DbtInteropCommand, DbtInvoker +from sqlbuild.integrations.dbt.types import DbtInvoker from sqlbuild.lint.models import ( CompileFacts, ExpandedLintFile, diff --git a/tests/unit/src/sqlbuild/compiler/pipeline/main/test_compile_phase_timings.py b/tests/unit/src/sqlbuild/compiler/pipeline/main/test_compile_phase_timings.py index e5dba7c66c..b17f2369ce 100644 --- a/tests/unit/src/sqlbuild/compiler/pipeline/main/test_compile_phase_timings.py +++ b/tests/unit/src/sqlbuild/compiler/pipeline/main/test_compile_phase_timings.py @@ -15,8 +15,6 @@ from sqlbuild.compiler.pipeline._helpers import materializations from sqlbuild.compiler.pipeline.models import CompilePipelineOptions, CompilePipelineResult from sqlbuild.compiler.planner.models import PlanOutput -from sqlbuild.diagnostics.models import PartialBuildPhaseTimings -from sqlbuild.rule_engine._helpers.run import rule_threads from sqlbuild.rule_engine.models import RulesRunResult from sqlbuild.runtime.contracts.models import ConnectionHooks from tests.unit.src.sqlbuild.compiler.pipeline.main._test_types import PipelinePhaseTimingTestCase diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/helpers.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/helpers.py index 535c17cbe0..26da58fa3c 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/helpers.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/helpers.py @@ -10,7 +10,6 @@ import pytest import sqlbuild._native as native -from sqlbuild import _native from sqlbuild._native import ( NativeRulesRequest, build_rule_catalogue, @@ -30,10 +29,7 @@ from sqlbuild.compiler.compile.types import CompiledProjectView, CompiledResourceType, SqlTestMode from sqlbuild.compiler.discovery.models import DiscoveredSqlTestBlock, DiscoveredSqlTestFile from sqlbuild.compiler.scopes.models import ScopeIndex -from sqlbuild.rule_engine._helpers.engine import custom_dispatch as native -from sqlbuild.rule_engine._helpers.engine.custom_rules import evaluate_custom_rules_cached -from sqlbuild.rule_engine.constants import MIN_CUSTOM_RULE_TEST_CASES, RULES_SKILL_PATHS -from sqlbuild.rule_engine.main._evaluate import evaluate +from sqlbuild.rule_engine.constants import MIN_CUSTOM_RULE_TEST_CASES from sqlbuild.rule_engine.models import ( CustomRulesOutcome, Finding, @@ -57,8 +53,8 @@ _MODEL_SQL: str = "WITH final AS (SELECT 1 AS id) SELECT id FROM final" -_BUILD_RULES_REQUEST: Callable[..., NativeRulesRequest] = _native.build_rules_request -_SCOPE_PROJECTION: Callable[..., dict[str, Any]] = _native.scope_metadata_projection +_BUILD_RULES_REQUEST: Callable[..., NativeRulesRequest] = native.build_rules_request +_SCOPE_PROJECTION: Callable[..., dict[str, Any]] = native.scope_metadata_projection def project_with_scope(*, index: ScopeIndex) -> CompiledProjectView: @@ -101,7 +97,7 @@ def evaluate_rules_request(request: object) -> NativeRulesEvaluation: reused=False, ) - monkeypatch.setattr(_native, "scope_metadata_projection", scope_metadata_projection) + monkeypatch.setattr(native, "scope_metadata_projection", scope_metadata_projection) monkeypatch.setattr(native, "build_rules_request", build_rules_request) monkeypatch.setattr(native, "evaluate_rules_request", evaluate_rules_request) native.evaluate_native_rules( From 10f4f97c14684b19b22b0c08f376da04037ea261 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 15 Oct 2026 09:19:58 +0100 Subject: [PATCH 08/12] refactor: host compute log storage natively Move LocalFilesystemComputeLogStorage, metadata encoding and validation, inventory and pruning into sqlbuild-cli's compute log module and expose the storage as a native class. The replay path and the storage now share one metadata encoder, inventory and prune implementation. Python keeps the stream tees, the scoped capture and the public shapes. --- .../src/compute_logs/_helpers/document.rs | 200 ++++++ .../src/compute_logs/_helpers/files.rs | 34 +- .../src/compute_logs/_helpers/metadata.rs | 150 +---- .../src/compute_logs/_helpers/mod.rs | 3 + .../src/compute_logs/_helpers/paths.rs | 160 +++++ .../src/compute_logs/_helpers/prune.rs | 145 ++--- .../src/compute_logs/_helpers/stored.rs | 117 ++++ .../src/compute_logs/_helpers/timestamps.rs | 61 +- .../src/compute_logs/constants.rs | 6 + .../sqlbuild-cli/src/compute_logs/errors.rs | 36 + .../compute_logs/main/capture_inventory.rs | 12 + .../compute_logs/main/capture_is_complete.rs | 14 + .../src/compute_logs/main/check_completion.rs | 59 ++ .../src/compute_logs/main/complete_capture.rs | 23 + .../src/compute_logs/main/delete_capture.rs | 15 + .../sqlbuild-cli/src/compute_logs/main/mod.rs | 12 + .../src/compute_logs/main/open_log_root.rs | 47 ++ .../src/compute_logs/main/prune_captures.rs | 72 ++ .../src/compute_logs/main/publish_capture.rs | 42 ++ .../main/read_capture_document.rs | 30 + .../compute_logs/main/read_stream_chunk.rs | 52 ++ .../src/compute_logs/main/start_capture.rs | 54 +- .../compute_logs/main/stream_byte_counts.rs | 21 + .../compute_logs/main/validate_capture_id.rs | 9 + crates/sqlbuild-cli/src/compute_logs/mod.rs | 1 + .../sqlbuild-cli/src/compute_logs/models.rs | 58 ++ crates/sqlbuild-cli/src/compute_logs/tests.rs | 2 + .../src/compute_logs/tests/test_metadata.rs | 12 +- .../compute_logs/tests/test_publication.rs | 48 ++ .../src/compute_logs/tests/test_types.rs | 9 + .../src/bindings/_helpers/cli/registration.rs | 11 +- .../_helpers/compute_logs/conversion.rs | 252 +++++++ .../src/bindings/_helpers/compute_logs/mod.rs | 3 + .../_helpers/compute_logs/registration.rs | 9 + .../bindings/_helpers/compute_logs/storage.rs | 431 ++++++++++++ .../src/bindings/_helpers/mod.rs | 1 + python-boundary.toml | 41 +- src/sqlbuild/_native.pyi | 56 +- .../entrypoint/_dispatch_with_compute_logs.py | 15 +- src/sqlbuild/compute_logs.py | 9 +- src/sqlbuild/runtime/compute_logs/README.md | 5 +- .../runtime/compute_logs/_helpers/metadata.py | 76 --- .../local_filesystem_compute_log_storage.py | 616 ------------------ .../classes/scoped_compute_log_capture.py | 4 +- .../runtime/compute_logs/main/__init__.py | 2 +- .../main/open_local_compute_log_storage.py | 23 - .../main/run_with_compute_log_capture.py | 29 + .../runtime/compute_logs/classes/helpers.py | 30 +- .../test_scoped_compute_log_capture.py | 17 +- .../test_storage_recovery_and_paths.py | 97 --- 50 files changed, 2070 insertions(+), 1161 deletions(-) create mode 100644 crates/sqlbuild-cli/src/compute_logs/_helpers/document.rs create mode 100644 crates/sqlbuild-cli/src/compute_logs/_helpers/paths.rs create mode 100644 crates/sqlbuild-cli/src/compute_logs/_helpers/stored.rs create mode 100644 crates/sqlbuild-cli/src/compute_logs/errors.rs create mode 100644 crates/sqlbuild-cli/src/compute_logs/main/capture_inventory.rs create mode 100644 crates/sqlbuild-cli/src/compute_logs/main/capture_is_complete.rs create mode 100644 crates/sqlbuild-cli/src/compute_logs/main/check_completion.rs create mode 100644 crates/sqlbuild-cli/src/compute_logs/main/complete_capture.rs create mode 100644 crates/sqlbuild-cli/src/compute_logs/main/delete_capture.rs create mode 100644 crates/sqlbuild-cli/src/compute_logs/main/open_log_root.rs create mode 100644 crates/sqlbuild-cli/src/compute_logs/main/prune_captures.rs create mode 100644 crates/sqlbuild-cli/src/compute_logs/main/publish_capture.rs create mode 100644 crates/sqlbuild-cli/src/compute_logs/main/read_capture_document.rs create mode 100644 crates/sqlbuild-cli/src/compute_logs/main/read_stream_chunk.rs create mode 100644 crates/sqlbuild-cli/src/compute_logs/main/stream_byte_counts.rs create mode 100644 crates/sqlbuild-cli/src/compute_logs/main/validate_capture_id.rs create mode 100644 crates/sqlbuild-cli/src/compute_logs/tests/test_publication.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/compute_logs/conversion.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/compute_logs/mod.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/compute_logs/registration.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/compute_logs/storage.rs delete mode 100644 src/sqlbuild/runtime/compute_logs/_helpers/metadata.py delete mode 100644 src/sqlbuild/runtime/compute_logs/classes/local_filesystem_compute_log_storage.py delete mode 100644 src/sqlbuild/runtime/compute_logs/main/open_local_compute_log_storage.py create mode 100644 src/sqlbuild/runtime/compute_logs/main/run_with_compute_log_capture.py diff --git a/crates/sqlbuild-cli/src/compute_logs/_helpers/document.rs b/crates/sqlbuild-cli/src/compute_logs/_helpers/document.rs new file mode 100644 index 0000000000..0c2266e1a7 --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/_helpers/document.rs @@ -0,0 +1,200 @@ +//! `metadata.json` encoding, decoding and the storage's metadata validation. + +use sqlbuild_core::json::main::dumps::dumps; +use sqlbuild_core::json::models::{JsonDialect, JsonInteger, JsonValue, StdlibJsonOptions}; + +use crate::compute_logs::_helpers::paths::validate_capture_id; +use crate::compute_logs::_helpers::timestamps::{parse_iso_instant, utc_date}; +use crate::compute_logs::constants::COMPUTE_LOG_FORMAT_VERSION; +use crate::compute_logs::errors::{ComputeLogError, ComputeLogErrorKind}; +use crate::compute_logs::models::{CaptureCompletion, CaptureDocument, CaptureInstant}; + +type Object = serde_json::Map; + +fn metadata_error(message: impl Into) -> ComputeLogError { + ComputeLogError::rejected(ComputeLogErrorKind::Metadata, message) +} + +fn text(value: &str) -> JsonValue { + JsonValue::String(value.to_owned()) +} + +fn integer(value: i64) -> JsonValue { + JsonValue::Integer(JsonInteger::from(value)) +} + +fn optional_text(value: Option<&str>) -> JsonValue { + value.map_or(JsonValue::Null, text) +} + +/// `metadata_to_json`: sorted keys, compact separators, non-ASCII kept, one trailing newline. +pub(crate) fn encode_document(document: &CaptureDocument) -> String { + let mut fields: Vec<(String, JsonValue)> = vec![ + ( + "format_version".to_owned(), + integer(document.format_version), + ), + ("invocation_id".to_owned(), text(&document.invocation_id)), + ("command".to_owned(), text(&document.command)), + ("project_dir".to_owned(), text(&document.project_dir)), + ("started_at".to_owned(), text(&document.started_at.text)), + ("capture_date".to_owned(), text(&document.capture_date)), + ( + "target".to_owned(), + optional_text(document.target.as_deref()), + ), + ( + "run_id".to_owned(), + optional_text(document.run_id.as_deref()), + ), + ]; + if let Some(completion) = &document.completion { + fields.extend([ + ( + "completed_at".to_owned(), + text(&completion.completed_at.text), + ), + ("exit_code".to_owned(), integer(completion.exit_code)), + ("stdout_bytes".to_owned(), integer(completion.stdout_bytes)), + ("stderr_bytes".to_owned(), integer(completion.stderr_bytes)), + ( + "diagnostics_bytes".to_owned(), + integer(completion.diagnostics_bytes), + ), + ]); + } + let options = StdlibJsonOptions { + ensure_ascii: false, + sort_keys: true, + ..StdlibJsonOptions::new(None).with_separators(",", ":") + }; + let encoded = + dumps(&JsonValue::Object(fields), &JsonDialect::Stdlib(options)).unwrap_or_default(); + format!("{encoded}\n") +} + +fn required<'a>(object: &'a Object, key: &str) -> Result<&'a serde_json::Value, ComputeLogError> { + object + .get(key) + .ok_or_else(|| metadata_error(format!("invalid capture metadata: '{key}'"))) +} + +fn required_text(object: &Object, key: &str) -> Result { + required(object, key)? + .as_str() + .map(str::to_owned) + .ok_or_else(|| metadata_error(format!("invalid capture metadata: {key} must be a string"))) +} + +fn optional_field(object: &Object, key: &str) -> Result, ComputeLogError> { + match object.get(key) { + None | Some(serde_json::Value::Null) => Ok(None), + Some(serde_json::Value::String(value)) => Ok(Some(value.clone())), + Some(_) => Err(metadata_error(format!( + "invalid capture metadata: {key} must be a string or null" + ))), + } +} + +fn instant(object: &Object, key: &str) -> Result { + let stored = required_text(object, key)?; + let (micros, aware) = parse_iso_instant(&stored.replace('Z', "+00:00")).ok_or_else(|| { + metadata_error(format!( + "invalid capture metadata: Invalid isoformat string: '{stored}'" + )) + })?; + Ok(CaptureInstant { + micros, + aware, + text: stored, + }) +} + +fn count(object: &Object, key: &str) -> Result { + let value = required(object, key)?; + if value.is_boolean() { + return Err(metadata_error("exit_code and byte counts must be integers")); + } + value + .as_i64() + .ok_or_else(|| metadata_error("exit_code and byte counts must be integers")) +} + +/// `metadata_from_json`: an initial or final capture document, or the metadata error raised. +pub(crate) fn decode_document(raw_json: &str) -> Result { + let parsed: serde_json::Value = serde_json::from_str(raw_json) + .map_err(|error| metadata_error(format!("invalid capture metadata: {error}")))?; + let serde_json::Value::Object(object) = parsed else { + return Err(metadata_error("capture metadata must be a JSON object")); + }; + let format_version = required(&object, "format_version")?; + let invocation_id = required_text(&object, "invocation_id")?; + let command = required_text(&object, "command")?; + let project_dir = required_text(&object, "project_dir")?; + let started_at = instant(&object, "started_at")?; + let capture_date = required_text(&object, "capture_date")?; + let target = optional_field(&object, "target")?; + let run_id = optional_field(&object, "run_id")?; + if format_version.is_boolean() || format_version.as_i64() != Some(COMPUTE_LOG_FORMAT_VERSION) { + return Err(metadata_error("unsupported compute log format version")); + } + let completion = if object.contains_key("completed_at") { + Some(CaptureCompletion { + completed_at: instant(&object, "completed_at")?, + exit_code: count(&object, "exit_code")?, + stdout_bytes: count(&object, "stdout_bytes")?, + stderr_bytes: count(&object, "stderr_bytes")?, + diagnostics_bytes: count(&object, "diagnostics_bytes")?, + }) + } else { + None + }; + Ok(CaptureDocument { + format_version: COMPUTE_LOG_FORMAT_VERSION, + invocation_id, + command, + project_dir, + started_at, + capture_date, + target, + run_id, + completion, + }) +} + +/// `_validate_metadata`: identity, format version, required text, timezones and byte counts. +pub(crate) fn validate_document(document: &CaptureDocument) -> Result<(), ComputeLogError> { + validate_capture_id(&document.invocation_id)?; + if document.format_version != COMPUTE_LOG_FORMAT_VERSION { + return Err(metadata_error("unsupported compute log format version")); + } + if document.command.trim().is_empty() || document.project_dir.trim().is_empty() { + return Err(metadata_error("command and project_dir must be non-empty")); + } + if !document.started_at.aware { + return Err(metadata_error("started_at must be timezone-aware")); + } + if document.capture_date != utc_date(document.started_at.micros) { + return Err(metadata_error( + "capture_date must derive from UTC started_at", + )); + } + if let Some(completion) = &document.completion { + if !completion.completed_at.aware + || completion.completed_at.micros < document.started_at.micros + { + return Err(metadata_error( + "completed_at must be aware and not precede started_at", + )); + } + if completion + .stdout_bytes + .min(completion.stderr_bytes) + .min(completion.diagnostics_bytes) + < 0 + { + return Err(metadata_error("final byte counts must be nonnegative")); + } + } + Ok(()) +} diff --git a/crates/sqlbuild-cli/src/compute_logs/_helpers/files.rs b/crates/sqlbuild-cli/src/compute_logs/_helpers/files.rs index 085ffc9abe..3d001759f7 100644 --- a/crates/sqlbuild-cli/src/compute_logs/_helpers/files.rs +++ b/crates/sqlbuild-cli/src/compute_logs/_helpers/files.rs @@ -4,7 +4,7 @@ use std::fs::{File, OpenOptions}; use std::io::{self, Write}; use std::path::Path; -use crate::compute_logs::constants::METADATA_FILE_NAME; +use crate::compute_logs::constants::{METADATA_FILE_NAME, STREAM_FILE_NAMES}; /// Random 128-bit hex text, as `uuid.uuid4().hex`. pub(crate) fn uuid4_hex() -> io::Result { @@ -49,3 +49,35 @@ fn write_and_replace(temporary: &Path, destination: &Path, json: &str) -> io::Re pub(crate) fn is_symlink(path: &Path) -> bool { std::fs::symlink_metadata(path).is_ok_and(|metadata| metadata.file_type().is_symlink()) } + +/// Publish `staging` as `directory` with its streams and metadata, removing it on failure. +pub(crate) fn publish_staged( + staging: &Path, + directory: &Path, + metadata_json: &str, +) -> io::Result<[File; 3]> { + let mut published: bool = false; + let result = (|| -> io::Result<[File; 3]> { + let mut streams: [File; 3] = [ + create_new(&staging.join(STREAM_FILE_NAMES[0]))?, + create_new(&staging.join(STREAM_FILE_NAMES[1]))?, + create_new(&staging.join(STREAM_FILE_NAMES[2]))?, + ]; + write_metadata(staging, metadata_json)?; + for stream in &mut streams { + stream.flush()?; + stream.sync_all()?; + } + sync_directory(staging)?; + std::fs::rename(staging, directory)?; + published = true; + if let Some(parent) = directory.parent() { + sync_directory(parent)?; + } + Ok(streams) + })(); + if result.is_err() { + let _ = std::fs::remove_dir_all(if published { directory } else { staging }); + } + result +} diff --git a/crates/sqlbuild-cli/src/compute_logs/_helpers/metadata.rs b/crates/sqlbuild-cli/src/compute_logs/_helpers/metadata.rs index 1ed3c70ad6..c528ba3013 100644 --- a/crates/sqlbuild-cli/src/compute_logs/_helpers/metadata.rs +++ b/crates/sqlbuild-cli/src/compute_logs/_helpers/metadata.rs @@ -1,36 +1,25 @@ -//! `metadata_to_json` and the subset of `metadata_from_json` validation pruning relies on. +//! `metadata_to_json` of a capture the native `sqb` replay records. -use sqlbuild_core::json::main::dumps::dumps; -use sqlbuild_core::json::models::{JsonDialect, JsonInteger, JsonValue, StdlibJsonOptions}; - -use crate::compute_logs::_helpers::timestamps::{ - epoch_microseconds, parse_utc_timestamp, utc_date, utc_timestamp, -}; +use crate::compute_logs::_helpers::document::encode_document; +use crate::compute_logs::_helpers::timestamps::{epoch_microseconds, utc_date, utc_timestamp}; use crate::compute_logs::constants::COMPUTE_LOG_FORMAT_VERSION; -use crate::compute_logs::models::{CaptureEnd, CaptureStart}; +use crate::compute_logs::models::{ + CaptureCompletion, CaptureDocument, CaptureEnd, CaptureInstant, CaptureStart, +}; /// The streams' final byte counts: stdout, stderr and diagnostics. pub(crate) type ByteCounts = (usize, usize, usize); -/// A stored capture's identity and the instant pruning orders it by. -#[derive(Clone, Debug, PartialEq, Eq)] -pub(crate) struct StoredCapture { - pub(crate) invocation_id: String, - pub(crate) capture_date: String, - /// `completed_at`, or `started_at` while incomplete, in microseconds. - pub(crate) order_instant: i64, -} - -fn text(value: &str) -> JsonValue { - JsonValue::String(value.to_owned()) -} - -fn integer(value: i64) -> JsonValue { - JsonValue::Integer(JsonInteger::from(value)) +fn utc_instant(micros: i64) -> CaptureInstant { + CaptureInstant { + micros, + aware: true, + text: utc_timestamp(micros), + } } -fn count(value: usize) -> JsonValue { - integer(i64::try_from(value).unwrap_or(i64::MAX)) +fn count(value: usize) -> i64 { + i64::try_from(value).unwrap_or(i64::MAX) } /// `metadata_to_json`: sorted keys, compact separators, non-ASCII kept, one trailing newline. @@ -39,98 +28,25 @@ pub(crate) fn metadata_json( end: Option<(&CaptureEnd, i64, ByteCounts)>, ) -> String { let started_at = epoch_microseconds(start.started_at); - let mut fields: Vec<(String, JsonValue)> = vec![ - ( - "format_version".to_owned(), - integer(COMPUTE_LOG_FORMAT_VERSION), - ), - ("invocation_id".to_owned(), text(invocation_id)), - ("command".to_owned(), text(&start.command)), - ("project_dir".to_owned(), text(&start.project_dir)), - ("started_at".to_owned(), text(&utc_timestamp(started_at))), - ("capture_date".to_owned(), text(&utc_date(started_at))), - ( - "target".to_owned(), - start.target.as_deref().map_or(JsonValue::Null, text), - ), - ("run_id".to_owned(), JsonValue::Null), - ]; - if let Some((end, completed_at, (stdout_bytes, stderr_bytes, diagnostics_bytes))) = end { - fields.extend([ - ( - "completed_at".to_owned(), - text(&utc_timestamp(completed_at)), - ), - ("exit_code".to_owned(), integer(end.exit_code)), - ("stdout_bytes".to_owned(), count(stdout_bytes)), - ("stderr_bytes".to_owned(), count(stderr_bytes)), - ("diagnostics_bytes".to_owned(), count(diagnostics_bytes)), - ]); - } - let options = StdlibJsonOptions { - ensure_ascii: false, - sort_keys: true, - ..StdlibJsonOptions::new(None).with_separators(",", ":") - }; - let encoded = - dumps(&JsonValue::Object(fields), &JsonDialect::Stdlib(options)).unwrap_or_default(); - format!("{encoded}\n") -} - -fn string_field<'a>( - object: &'a serde_json::Map, - key: &str, -) -> Option<&'a str> { - object.get(key).and_then(serde_json::Value::as_str) -} - -/// The capture a stored `metadata.json` describes, or `None` where Python's validation rejects it. -pub(crate) fn stored_capture(json: &str) -> Option { - let Ok(serde_json::Value::Object(object)) = serde_json::from_str::(json) - else { - return None; - }; - if object - .get("format_version") - .and_then(serde_json::Value::as_i64) - != Some(COMPUTE_LOG_FORMAT_VERSION) - { - return None; - } - let invocation_id = string_field(&object, "invocation_id")?; - let capture_date = string_field(&object, "capture_date")?; - let command = string_field(&object, "command")?; - let project_dir = string_field(&object, "project_dir")?; - let started_at = parse_utc_timestamp(string_field(&object, "started_at")?)?; - if command.trim().is_empty() - || project_dir.trim().is_empty() - || utc_date(started_at) != capture_date - { - return None; - } - let order_instant = match object.get("completed_at") { - None => started_at, - Some(value) => { - let completed_at = parse_utc_timestamp(value.as_str()?)?; - let integers = [ - "exit_code", - "stdout_bytes", - "stderr_bytes", - "diagnostics_bytes", - ]; - if completed_at < started_at - || !integers - .iter() - .all(|key| object.get(*key).is_some_and(serde_json::Value::is_i64)) - { - return None; - } - completed_at - } - }; - Some(StoredCapture { + encode_document(&CaptureDocument { + format_version: COMPUTE_LOG_FORMAT_VERSION, invocation_id: invocation_id.to_owned(), - capture_date: capture_date.to_owned(), - order_instant, + command: start.command.clone(), + project_dir: start.project_dir.clone(), + started_at: utc_instant(started_at), + capture_date: utc_date(started_at), + target: start.target.clone(), + run_id: None, + completion: end.map( + |(end, completed_at, (stdout_bytes, stderr_bytes, diagnostics_bytes))| { + CaptureCompletion { + completed_at: utc_instant(completed_at), + exit_code: end.exit_code, + stdout_bytes: count(stdout_bytes), + stderr_bytes: count(stderr_bytes), + diagnostics_bytes: count(diagnostics_bytes), + } + }, + ), }) } diff --git a/crates/sqlbuild-cli/src/compute_logs/_helpers/mod.rs b/crates/sqlbuild-cli/src/compute_logs/_helpers/mod.rs index cecf8b0763..0201d6838f 100644 --- a/crates/sqlbuild-cli/src/compute_logs/_helpers/mod.rs +++ b/crates/sqlbuild-cli/src/compute_logs/_helpers/mod.rs @@ -1,4 +1,7 @@ +pub(crate) mod document; pub(crate) mod files; pub(crate) mod metadata; +pub(crate) mod paths; pub(crate) mod prune; +pub(crate) mod stored; pub(crate) mod timestamps; diff --git a/crates/sqlbuild-cli/src/compute_logs/_helpers/paths.rs b/crates/sqlbuild-cli/src/compute_logs/_helpers/paths.rs new file mode 100644 index 0000000000..5f42964c55 --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/_helpers/paths.rs @@ -0,0 +1,160 @@ +//! Capture identities, date directory names, symlink rejection and root containment. + +use std::path::{Path, PathBuf}; + +use crate::compute_logs::_helpers::files::is_symlink; +use crate::compute_logs::constants::{ + CAPTURE_DATE_HYPHENS, CAPTURE_DATE_LENGTH, MAX_CAPTURE_ID_LENGTH, MAX_READ_BYTES, +}; +use crate::compute_logs::errors::{ComputeLogError, ComputeLogErrorKind}; + +/// Python's `[A-Za-z0-9][A-Za-z0-9_-]{0,127}\Z` capture identity. +pub(crate) fn validate_capture_id(invocation_id: &str) -> Result<(), ComputeLogError> { + let bytes = invocation_id.as_bytes(); + let valid = !bytes.is_empty() + && bytes.len() <= MAX_CAPTURE_ID_LENGTH + && bytes[0].is_ascii_alphanumeric() + && bytes + .iter() + .all(|byte| byte.is_ascii_alphanumeric() || *byte == b'_' || *byte == b'-'); + if valid { + return Ok(()); + } + Err(ComputeLogError::rejected( + ComputeLogErrorKind::InvalidCaptureId, + "invocation_id must contain only ASCII letters, digits, underscores, or hyphens", + )) +} + +/// Python's `\d{4}-\d{2}-\d{2}\Z` date directory name. +pub(crate) fn is_capture_date(name: &str) -> bool { + let bytes = name.as_bytes(); + bytes.len() == CAPTURE_DATE_LENGTH + && bytes.iter().enumerate().all(|(index, byte)| { + if CAPTURE_DATE_HYPHENS.contains(&index) { + *byte == b'-' + } else { + byte.is_ascii_digit() + } + }) +} + +/// The final component of `path` as text. +pub(crate) fn file_name(path: &Path) -> String { + path.file_name() + .map(|name| name.to_string_lossy().into_owned()) + .unwrap_or_default() +} + +/// `_ensure_not_symlink`. +pub(crate) fn ensure_not_symlink(path: &Path) -> Result<(), ComputeLogError> { + if is_symlink(path) { + return Err(ComputeLogError::rejected( + ComputeLogErrorKind::Path, + format!("compute log path cannot be a symlink: {}", path.display()), + )); + } + Ok(()) +} + +/// `Path.resolve(strict=False)`: the deepest existing ancestor resolved, the rest appended. +fn resolve_lenient(path: &Path) -> PathBuf { + let mut missing: Vec = Vec::new(); + let mut existing: &Path = path; + loop { + if let Ok(resolved) = std::fs::canonicalize(existing) { + return missing + .iter() + .rev() + .fold(resolved, |resolved, name| resolved.join(name)); + } + match (existing.parent(), existing.file_name()) { + (Some(parent), Some(name)) => { + missing.push(name.to_owned()); + existing = parent; + } + _ => return path.to_path_buf(), + } + } +} + +/// `_ensure_contained`: `path` resolves beneath the configured `root`. +pub(crate) fn ensure_contained(root: &Path, path: &Path) -> Result<(), ComputeLogError> { + if resolve_lenient(path).starts_with(root) { + return Ok(()); + } + Err(ComputeLogError::rejected( + ComputeLogErrorKind::Path, + format!( + "compute log path escapes configured root: {}", + path.display() + ), + )) +} + +/// `_find_capture`: the one capture directory named `invocation_id` under any date directory. +pub(crate) fn find_capture( + root: &Path, + invocation_id: &str, +) -> Result, ComputeLogError> { + validate_capture_id(invocation_id)?; + let mut matches: Vec = Vec::new(); + for entry in std::fs::read_dir(root)? { + let date_directory = entry?.path(); + if !is_capture_date(&file_name(&date_directory)) { + continue; + } + ensure_not_symlink(&date_directory)?; + ensure_contained(root, &date_directory)?; + if !date_directory.is_dir() { + continue; + } + let candidate = date_directory.join(invocation_id); + ensure_not_symlink(&candidate)?; + if candidate.exists() { + ensure_contained(root, &candidate)?; + matches.push(candidate); + } + } + if matches.len() > 1 { + return Err(ComputeLogError::rejected( + ComputeLogErrorKind::Metadata, + format!("duplicate capture identity '{invocation_id}'"), + )); + } + Ok(matches.pop()) +} + +/// `_find_capture` without `missing_ok`. +pub(crate) fn existing_capture( + root: &Path, + invocation_id: &str, +) -> Result { + find_capture(root, invocation_id)?.ok_or_else(|| { + ComputeLogError::rejected( + ComputeLogErrorKind::NotFound, + format!("capture not found for invocation '{invocation_id}'"), + ) + }) +} + +/// A read's nonnegative cursor and its byte limit, positive and no greater than the maximum. +pub(crate) fn read_bounds(cursor: i64, max_bytes: i64) -> Result<(u64, u64), ComputeLogError> { + let cursor = u64::try_from(cursor).map_err(|_| { + ComputeLogError::rejected( + ComputeLogErrorKind::InvalidCursor, + "cursor must be a nonnegative integer", + ) + })?; + let limit_error = || { + ComputeLogError::rejected( + ComputeLogErrorKind::InvalidLimit, + format!("max_bytes must be a positive integer no greater than {MAX_READ_BYTES}"), + ) + }; + let limit = u64::try_from(max_bytes).map_err(|_| limit_error())?; + if !(1..=MAX_READ_BYTES).contains(&limit) { + return Err(limit_error()); + } + Ok((cursor, limit)) +} diff --git a/crates/sqlbuild-cli/src/compute_logs/_helpers/prune.rs b/crates/sqlbuild-cli/src/compute_logs/_helpers/prune.rs index 33e47e0194..f56f91540d 100644 --- a/crates/sqlbuild-cli/src/compute_logs/_helpers/prune.rs +++ b/crates/sqlbuild-cli/src/compute_logs/_helpers/prune.rs @@ -1,125 +1,62 @@ -//! `LocalFilesystemComputeLogStorage.prune`: keep the newest complete captures under a file lock. +//! The prune file lock Python's `filelock` takes, and the best-effort prune after a replay. -use std::fs::OpenOptions; -use std::path::{Path, PathBuf}; +use std::fs::{File, OpenOptions}; +use std::path::Path; -use crate::compute_logs::_helpers::files::is_symlink; -use crate::compute_logs::_helpers::metadata::{StoredCapture, stored_capture}; -use crate::compute_logs::constants::{ - CAPTURE_DATE_HYPHENS, CAPTURE_DATE_LENGTH, COMPLETE_FILE_NAME, METADATA_FILE_NAME, - PRUNE_LOCK_FILE_NAME, -}; +use crate::compute_logs::constants::PRUNE_LOCK_FILE_NAME; +use crate::compute_logs::main::prune_captures::prune_captures; -/// Python's `\d{4}-\d{2}-\d{2}\Z` date directory name. -fn is_capture_date(name: &str) -> bool { - let bytes = name.as_bytes(); - bytes.len() == CAPTURE_DATE_LENGTH - && bytes.iter().enumerate().all(|(index, byte)| { - if CAPTURE_DATE_HYPHENS.contains(&index) { - *byte == b'-' - } else { - byte.is_ascii_digit() - } - }) +/// Open or create `/.prune.lock`. +fn open_lock(root: &Path) -> std::io::Result { + OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(false) + .open(root.join(PRUNE_LOCK_FILE_NAME)) } -fn sorted_entries(directory: &Path) -> Option> { - let Ok(listing) = std::fs::read_dir(directory) else { - return None; - }; - let mut entries: Vec = Vec::new(); - for entry in listing { - let Ok(entry) = entry else { - return None; - }; - entries.push(entry.path()); - } - entries.sort(); - Some(entries) +/// Python `filelock`'s `fcntl.flock` exclusive lock, released when the file closes. +#[cfg(unix)] +fn lock_exclusive(file: &File) -> std::io::Result<()> { + rustix::fs::flock(file, rustix::fs::FlockOperation::LockExclusive).map_err(std::io::Error::from) } -fn file_name(path: &Path) -> String { - path.file_name() - .map(|name| name.to_string_lossy().into_owned()) - .unwrap_or_default() +/// Captures are never published here, as in Python, whose directory `fsync` fails on Windows. +#[cfg(not(unix))] +fn lock_exclusive(_file: &File) -> std::io::Result<()> { + Err(std::io::Error::from(std::io::ErrorKind::Unsupported)) } -/// Every complete capture as Python's `inventory` lists it, or `None` where it raises. -fn complete_captures(root: &Path) -> Option> { - let mut complete: Vec<(StoredCapture, PathBuf)> = Vec::new(); - for date_directory in sorted_entries(root)? { - if !is_capture_date(&file_name(&date_directory)) { - continue; - } - if is_symlink(&date_directory) { - return None; - } - if !date_directory.is_dir() { - continue; - } - for capture_directory in sorted_entries(&date_directory)? { - let name = file_name(&capture_directory); - if name.starts_with('.') { - continue; - } - if is_symlink(&capture_directory) { - return None; - } - if !capture_directory.is_dir() { - continue; - } - let Ok(json) = std::fs::read_to_string(capture_directory.join(METADATA_FILE_NAME)) - else { - return None; - }; - let capture = stored_capture(&json)?; - if capture.invocation_id != name || capture.capture_date != file_name(&date_directory) { - return None; - } - if capture_directory.join(COMPLETE_FILE_NAME).is_file() { - complete.push((capture, capture_directory)); - } +/// Hold the exclusive prune lock, retrying when a releasing holder unlinked the file first. +pub(crate) fn acquire_lock(root: &Path) -> std::io::Result { + loop { + let file = open_lock(root)?; + lock_exclusive(&file)?; + if linked(&file)? { + return Ok(file); } } - Some(complete) } -/// Python `filelock`'s `fcntl.flock` exclusive lock, released when the file closes. #[cfg(unix)] -fn lock_exclusive(file: &std::fs::File) -> std::io::Result<()> { - rustix::fs::flock(file, rustix::fs::FlockOperation::LockExclusive).map_err(std::io::Error::from) +fn linked(file: &File) -> std::io::Result { + use std::os::unix::fs::MetadataExt; + Ok(file.metadata()?.nlink() > 0) } -/// Captures are never published here, as in Python, whose directory `fsync` fails on Windows. #[cfg(not(unix))] -fn lock_exclusive(_file: &std::fs::File) -> std::io::Result<()> { - Err(std::io::Error::from(std::io::ErrorKind::Unsupported)) +fn linked(_file: &File) -> std::io::Result { + Ok(true) } -/// Delete all but the `retain` most recently completed captures; any fault leaves them all. +/// Release the lock as `filelock` does: unlink the file, then close it. +pub(crate) fn release_lock(root: &Path, file: File) { + let _ = std::fs::remove_file(root.join(PRUNE_LOCK_FILE_NAME)); + drop(file); +} + +/// Delete all but the `retain` most recently completed captures; any fault is ignored. pub(crate) fn prune(root: &Path, retain: usize) { - let Ok(lock) = OpenOptions::new() - .read(true) - .write(true) - .create(true) - .truncate(false) - .open(root.join(PRUNE_LOCK_FILE_NAME)) - else { - return; - }; - if lock_exclusive(&lock).is_err() { - return; - } - if let Some(mut complete) = complete_captures(root) { - complete.sort_by(|left, right| { - (right.0.order_instant, &right.0.invocation_id) - .cmp(&(left.0.order_instant, &left.0.invocation_id)) - }); - for (_, directory) in complete.iter().skip(retain) { - if directory.join(COMPLETE_FILE_NAME).is_file() { - let _ = std::fs::remove_dir_all(directory); - } - } - } - drop(lock); + let _ = prune_captures(root, Some(retain), &|_| false); } diff --git a/crates/sqlbuild-cli/src/compute_logs/_helpers/stored.rs b/crates/sqlbuild-cli/src/compute_logs/_helpers/stored.rs new file mode 100644 index 0000000000..89627276a9 --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/_helpers/stored.rs @@ -0,0 +1,117 @@ +//! Reading captures already on disk: directory identity, inventory and deletion. + +use std::path::{Path, PathBuf}; + +use crate::compute_logs::_helpers::document::{decode_document, validate_document}; +use crate::compute_logs::_helpers::paths::{ + ensure_contained, ensure_not_symlink, existing_capture, file_name, is_capture_date, +}; +use crate::compute_logs::constants::{ + COMPLETE_FILE_NAME, METADATA_FILE_NAME, PRUNE_LOCK_FILE_NAME, +}; +use crate::compute_logs::errors::{ComputeLogError, ComputeLogErrorKind}; +use crate::compute_logs::models::{CaptureDocument, InventoryEntry}; + +/// `_validate_directory_metadata`: the directory names are authoritative for identity and date. +pub(crate) fn validate_directory_document( + directory: &Path, + document: &CaptureDocument, +) -> Result<(), ComputeLogError> { + if document.invocation_id != file_name(directory) { + return Err(ComputeLogError::rejected( + ComputeLogErrorKind::Metadata, + "metadata invocation_id does not match capture directory", + )); + } + let date = directory.parent().map(file_name).unwrap_or_default(); + if document.capture_date != date { + return Err(ComputeLogError::rejected( + ComputeLogErrorKind::Metadata, + "metadata capture_date does not match date directory", + )); + } + Ok(()) +} + +fn sorted_entries(directory: &Path) -> Result, ComputeLogError> { + let mut entries: Vec = Vec::new(); + for entry in std::fs::read_dir(directory)? { + entries.push(entry?.path()); + } + entries.sort(); + Ok(entries) +} + +fn stored_entry(root: &Path, directory: PathBuf) -> Result { + let metadata_path = directory.join(METADATA_FILE_NAME); + let marker = directory.join(COMPLETE_FILE_NAME); + ensure_contained(root, &metadata_path)?; + ensure_contained(root, &marker)?; + let document = decode_document(&std::fs::read_to_string(&metadata_path)?)?; + validate_document(&document)?; + validate_directory_document(&directory, &document)?; + Ok(InventoryEntry { + document, + is_complete: marker.is_file(), + path: directory, + }) +} + +/// `inventory`: every capture under dated directories, ordered by date and invocation id. +pub(crate) fn inventory(root: &Path) -> Result, ComputeLogError> { + let mut captures: Vec = Vec::new(); + for date_directory in sorted_entries(root)? { + let date_name = file_name(&date_directory); + if date_name == PRUNE_LOCK_FILE_NAME || !is_capture_date(&date_name) { + continue; + } + ensure_not_symlink(&date_directory)?; + ensure_contained(root, &date_directory)?; + if !date_directory.is_dir() { + continue; + } + for directory in sorted_entries(&date_directory)? { + if file_name(&directory).starts_with('.') { + continue; + } + ensure_not_symlink(&directory)?; + ensure_contained(root, &directory)?; + if !directory.is_dir() { + continue; + } + captures.push(stored_entry(root, directory)?); + } + } + captures.sort_by(|left, right| { + (&left.document.capture_date, &left.document.invocation_id) + .cmp(&(&right.document.capture_date, &right.document.invocation_id)) + }); + Ok(captures) +} + +/// `delete`: remove a complete capture no storage instance is writing. +pub(crate) fn delete_inactive( + root: &Path, + invocation_id: &str, + active: bool, +) -> Result<(), ComputeLogError> { + crate::compute_logs::_helpers::paths::validate_capture_id(invocation_id)?; + if active { + return Err(ComputeLogError::rejected( + ComputeLogErrorKind::State, + "active captures cannot be deleted", + )); + } + let directory = existing_capture(root, invocation_id)?; + let marker = directory.join(COMPLETE_FILE_NAME); + ensure_contained(root, &marker)?; + if !marker.is_file() { + return Err(ComputeLogError::rejected( + ComputeLogErrorKind::State, + "incomplete captures cannot be deleted", + )); + } + ensure_contained(root, &directory)?; + std::fs::remove_dir_all(&directory)?; + Ok(()) +} diff --git a/crates/sqlbuild-cli/src/compute_logs/_helpers/timestamps.rs b/crates/sqlbuild-cli/src/compute_logs/_helpers/timestamps.rs index ea16ce042f..0bcb995724 100644 --- a/crates/sqlbuild-cli/src/compute_logs/_helpers/timestamps.rs +++ b/crates/sqlbuild-cli/src/compute_logs/_helpers/timestamps.rs @@ -3,8 +3,9 @@ use std::time::{SystemTime, UNIX_EPOCH}; use crate::compute_logs::constants::{ - DAYS_PER_ERA, EPOCH_DAY_OFFSET, FEBRUARY, MARCH_BASED_JANUARY, MICROSECOND_DIGITS, - MICROSECONDS_PER_SECOND, SECONDS_PER_DAY, SECONDS_PER_HOUR, SECONDS_PER_MINUTE, YEARS_PER_ERA, + CAPTURE_DATE_LENGTH, DAYS_PER_ERA, EPOCH_DAY_OFFSET, FEBRUARY, HOURS_PER_DAY, + MARCH_BASED_JANUARY, MAX_DAY_OF_MONTH, MICROSECOND_DIGITS, MICROSECONDS_PER_SECOND, + MONTHS_PER_YEAR, SECONDS_PER_DAY, SECONDS_PER_HOUR, SECONDS_PER_MINUTE, YEARS_PER_ERA, }; /// Microseconds since the Unix epoch of `instant`. @@ -79,25 +80,59 @@ fn number(text: &str) -> Option { Some(value) } -/// Microseconds of a `YYYY-MM-DDTHH:MM:SS[.ffffff]` UTC timestamp ending in `Z` or `+00:00`. -pub(crate) fn parse_utc_timestamp(text: &str) -> Option { - let body = text - .strip_suffix('Z') - .or_else(|| text.strip_suffix("+00:00"))?; - let (date, time) = body.split_once('T')?; - let (clock, fraction) = time.split_once('.').unwrap_or((time, "")); +fn offset_seconds(text: &str) -> Option { + let (sign, body) = match text.as_bytes().first()? { + b'+' => (1, text.get(1..)?), + b'-' => (-1, text.get(1..)?), + _ => return None, + }; + let (whole, _fraction) = body.split_once('.').unwrap_or((body, "")); + let parts: Vec = whole.split(':').map(number).collect::>()?; + let seconds = match parts.as_slice() { + [hours, minutes] => hours * SECONDS_PER_HOUR + minutes * SECONDS_PER_MINUTE, + [hours, minutes, seconds] => { + hours * SECONDS_PER_HOUR + minutes * SECONDS_PER_MINUTE + seconds + } + _ => return None, + }; + Some(sign * seconds) +} + +/// UTC microseconds and awareness of `datetime.fromisoformat(text)`, or `None` when unparsable. +pub(crate) fn parse_iso_instant(text: &str) -> Option<(i64, bool)> { + let date = text.get(..CAPTURE_DATE_LENGTH)?; + let rest = text.get(CAPTURE_DATE_LENGTH + 1..)?; + let split = rest.find(['+', '-']).unwrap_or(rest.len()); + let (clock, offset) = rest.split_at(split); + let (whole, fraction) = clock.split_once('.').unwrap_or((clock, "")); let dates: Vec = date.split('-').map(number).collect::>()?; - let clocks: Vec = clock.split(':').map(number).collect::>()?; + let clocks: Vec = whole.split(':').map(number).collect::>()?; let ([year, month, day], [hour, minute, second]) = (dates.as_slice(), clocks.as_slice()) else { return None; }; + let in_range = (1..=MONTHS_PER_YEAR).contains(month) + && (1..=MAX_DAY_OF_MONTH).contains(day) + && (0..HOURS_PER_DAY).contains(hour) + && (0..SECONDS_PER_MINUTE).contains(minute) + && (0..SECONDS_PER_MINUTE).contains(second); + if !in_range || !fraction.bytes().all(|byte| byte.is_ascii_digit()) { + return None; + } let micros = match fraction { "" => 0, digits => number(format!("{digits:0<6}").get(..MICROSECOND_DIGITS)?)?, }; - let seconds = days_from_civil(*year, *month, *day) * SECONDS_PER_DAY + let local = (days_from_civil(*year, *month, *day) * SECONDS_PER_DAY + hour * SECONDS_PER_HOUR + minute * SECONDS_PER_MINUTE - + second; - Some(seconds * MICROSECONDS_PER_SECOND + micros) + + second) + * MICROSECONDS_PER_SECOND + + micros; + if offset.is_empty() { + return Some((local, false)); + } + Some(( + local - offset_seconds(offset)? * MICROSECONDS_PER_SECOND, + true, + )) } diff --git a/crates/sqlbuild-cli/src/compute_logs/constants.rs b/crates/sqlbuild-cli/src/compute_logs/constants.rs index c0ee017821..54650af368 100644 --- a/crates/sqlbuild-cli/src/compute_logs/constants.rs +++ b/crates/sqlbuild-cli/src/compute_logs/constants.rs @@ -27,3 +27,9 @@ pub const EPOCH_DAY_OFFSET: i64 = 719_468; /// `YYYY-MM-DD`: the length of a capture date directory name and its hyphen positions. pub const CAPTURE_DATE_LENGTH: usize = 10; pub const CAPTURE_DATE_HYPHENS: [usize; 2] = [4, 7]; +pub const MAX_READ_BYTES: u64 = 1_048_576; +pub const MAX_CAPTURE_ID_LENGTH: usize = 128; +pub const DEFAULT_LOG_DIRECTORY: &str = "logs"; +pub const MONTHS_PER_YEAR: i64 = 12; +pub const MAX_DAY_OF_MONTH: i64 = 31; +pub const HOURS_PER_DAY: i64 = 24; diff --git a/crates/sqlbuild-cli/src/compute_logs/errors.rs b/crates/sqlbuild-cli/src/compute_logs/errors.rs new file mode 100644 index 0000000000..2f6e9d16e6 --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/errors.rs @@ -0,0 +1,36 @@ +//! Compute log storage rejections, each naming the Python `ComputeLogStorageError` subclass. + +use std::io; + +/// Which `sqlbuild.runtime.compute_logs.exceptions` class a rejection raises. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ComputeLogErrorKind { + InvalidCaptureId, + InvalidCursor, + InvalidLimit, + AlreadyExists, + NotFound, + State, + Path, + Metadata, +} + +/// A storage rejection with its message, or an operating system failure Python lets propagate. +#[derive(Debug)] +pub enum ComputeLogError { + Rejected(ComputeLogErrorKind, String), + Io(io::Error), +} + +impl ComputeLogError { + /// A rejection of `kind` with `message`. + pub fn rejected(kind: ComputeLogErrorKind, message: impl Into) -> Self { + Self::Rejected(kind, message.into()) + } +} + +impl From for ComputeLogError { + fn from(error: io::Error) -> Self { + Self::Io(error) + } +} diff --git a/crates/sqlbuild-cli/src/compute_logs/main/capture_inventory.rs b/crates/sqlbuild-cli/src/compute_logs/main/capture_inventory.rs new file mode 100644 index 0000000000..e5afcc2771 --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/main/capture_inventory.rs @@ -0,0 +1,12 @@ +//! `inventory`: every complete and abandoned capture beneath the root. + +use std::path::Path; + +use crate::compute_logs::_helpers::stored::inventory; +use crate::compute_logs::errors::ComputeLogError; +use crate::compute_logs::models::InventoryEntry; + +/// Captures ordered by capture date and invocation id. +pub fn capture_inventory(root: &Path) -> Result, ComputeLogError> { + inventory(root) +} diff --git a/crates/sqlbuild-cli/src/compute_logs/main/capture_is_complete.rs b/crates/sqlbuild-cli/src/compute_logs/main/capture_is_complete.rs new file mode 100644 index 0000000000..c0b678b4e3 --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/main/capture_is_complete.rs @@ -0,0 +1,14 @@ +//! `is_complete`: whether a capture's completion marker exists. + +use std::path::Path; + +use crate::compute_logs::_helpers::paths::{ensure_contained, existing_capture}; +use crate::compute_logs::constants::COMPLETE_FILE_NAME; +use crate::compute_logs::errors::ComputeLogError; + +/// Capture-close evidence, without asserting the command's lifecycle status. +pub fn capture_is_complete(root: &Path, invocation_id: &str) -> Result { + let marker = existing_capture(root, invocation_id)?.join(COMPLETE_FILE_NAME); + ensure_contained(root, &marker)?; + Ok(marker.is_file()) +} diff --git a/crates/sqlbuild-cli/src/compute_logs/main/check_completion.rs b/crates/sqlbuild-cli/src/compute_logs/main/check_completion.rs new file mode 100644 index 0000000000..78966b850b --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/main/check_completion.rs @@ -0,0 +1,59 @@ +//! `mark_complete`: check final metadata against the capture before publishing it. + +use std::path::{Path, PathBuf}; + +use crate::compute_logs::_helpers::document::validate_document; +use crate::compute_logs::errors::{ComputeLogError, ComputeLogErrorKind}; +use crate::compute_logs::main::read_capture_document::read_capture_document; +use crate::compute_logs::models::CaptureDocument; + +fn metadata_error(message: &str) -> ComputeLogError { + ComputeLogError::rejected(ComputeLogErrorKind::Metadata, message) +} + +/// The capture directory once `document` matches its start and the open streams' `counts`. +pub fn check_completion( + root: &Path, + document: &CaptureDocument, + counts: Option<[i64; 3]>, +) -> Result { + validate_document(document)?; + let Some(counts) = counts else { + return Err(ComputeLogError::rejected( + ComputeLogErrorKind::State, + "capture is not active in this storage instance", + )); + }; + let mismatch = || metadata_error("final metadata does not match capture start metadata"); + let completion = document.completion.as_ref().ok_or_else(mismatch)?; + let (initial, directory) = read_capture_document(root, &document.invocation_id)?; + let same_start = ( + initial.started_at.micros, + &initial.capture_date, + &initial.command, + &initial.project_dir, + &initial.target, + &initial.run_id, + ) == ( + document.started_at.micros, + &document.capture_date, + &document.command, + &document.project_dir, + &document.target, + &document.run_id, + ); + if initial.completion.is_some() || !same_start { + return Err(mismatch()); + } + let expected = [ + completion.stdout_bytes, + completion.stderr_bytes, + completion.diagnostics_bytes, + ]; + if counts != expected { + return Err(metadata_error( + "final byte counts do not match captured streams", + )); + } + Ok(directory) +} diff --git a/crates/sqlbuild-cli/src/compute_logs/main/complete_capture.rs b/crates/sqlbuild-cli/src/compute_logs/main/complete_capture.rs new file mode 100644 index 0000000000..df010e96ca --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/main/complete_capture.rs @@ -0,0 +1,23 @@ +//! `mark_complete`: close a checked capture's streams and publish its final metadata and marker. + +use std::fs::File; +use std::path::Path; + +use crate::compute_logs::_helpers::document::encode_document; +use crate::compute_logs::_helpers::files::{create_new, sync_directory, write_metadata}; +use crate::compute_logs::constants::COMPLETE_FILE_NAME; +use crate::compute_logs::errors::ComputeLogError; +use crate::compute_logs::models::CaptureDocument; + +/// Close the already synced `streams`, then publish the final metadata and completion marker. +pub fn complete_capture( + directory: &Path, + document: &CaptureDocument, + streams: [File; 3], +) -> Result<(), ComputeLogError> { + drop(streams); + write_metadata(directory, &encode_document(document))?; + create_new(&directory.join(COMPLETE_FILE_NAME))?.sync_all()?; + sync_directory(directory)?; + Ok(()) +} diff --git a/crates/sqlbuild-cli/src/compute_logs/main/delete_capture.rs b/crates/sqlbuild-cli/src/compute_logs/main/delete_capture.rs new file mode 100644 index 0000000000..2cce863f06 --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/main/delete_capture.rs @@ -0,0 +1,15 @@ +//! `delete`: remove one complete, inactive capture. + +use std::path::Path; + +use crate::compute_logs::_helpers::stored::delete_inactive; +use crate::compute_logs::errors::ComputeLogError; + +/// Delete the capture; `active` captures and incomplete ones are refused. +pub fn delete_capture( + root: &Path, + invocation_id: &str, + active: bool, +) -> Result<(), ComputeLogError> { + delete_inactive(root, invocation_id, active) +} diff --git a/crates/sqlbuild-cli/src/compute_logs/main/mod.rs b/crates/sqlbuild-cli/src/compute_logs/main/mod.rs index 008c6b31d2..8e44c3eed2 100644 --- a/crates/sqlbuild-cli/src/compute_logs/main/mod.rs +++ b/crates/sqlbuild-cli/src/compute_logs/main/mod.rs @@ -1,2 +1,14 @@ +pub mod capture_inventory; +pub mod capture_is_complete; +pub mod check_completion; +pub mod complete_capture; +pub mod delete_capture; pub mod finish_capture; +pub mod open_log_root; +pub mod prune_captures; +pub mod publish_capture; +pub mod read_capture_document; +pub mod read_stream_chunk; pub mod start_capture; +pub mod stream_byte_counts; +pub mod validate_capture_id; diff --git a/crates/sqlbuild-cli/src/compute_logs/main/open_log_root.rs b/crates/sqlbuild-cli/src/compute_logs/main/open_log_root.rs new file mode 100644 index 0000000000..e3d8143965 --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/main/open_log_root.rs @@ -0,0 +1,47 @@ +//! Resolve and create the configured compute log root, as the storage constructor does. + +use std::path::{Path, PathBuf}; + +use crate::compute_logs::_helpers::files::is_symlink; +use crate::compute_logs::constants::DEFAULT_LOG_DIRECTORY; +use crate::compute_logs::errors::{ComputeLogError, ComputeLogErrorKind}; + +fn symlink_root() -> ComputeLogError { + ComputeLogError::rejected( + ComputeLogErrorKind::Path, + "configured compute log root cannot be a symlink", + ) +} + +/// The resolved root, `/logs` unless `root` is given; rejects symlinks and files. +pub fn open_log_root( + project_dir: &Path, + root: Option<&Path>, + retention_count: Option, +) -> Result { + if retention_count.is_some_and(|count| count < 0) { + return Err(ComputeLogError::rejected( + ComputeLogErrorKind::InvalidLimit, + "retention_count must be nonnegative or None", + )); + } + let configured: PathBuf = root.map_or_else( + || project_dir.join(DEFAULT_LOG_DIRECTORY), + Path::to_path_buf, + ); + if is_symlink(&configured) { + return Err(symlink_root()); + } + std::fs::create_dir_all(&configured)?; + if is_symlink(&configured) { + return Err(symlink_root()); + } + let resolved = std::fs::canonicalize(&configured)?; + if !resolved.is_dir() { + return Err(ComputeLogError::rejected( + ComputeLogErrorKind::Path, + "compute log root must be a directory", + )); + } + Ok(resolved) +} diff --git a/crates/sqlbuild-cli/src/compute_logs/main/prune_captures.rs b/crates/sqlbuild-cli/src/compute_logs/main/prune_captures.rs new file mode 100644 index 0000000000..5e9d8d9cdb --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/main/prune_captures.rs @@ -0,0 +1,72 @@ +//! `prune`: keep the newest complete captures under the prune file lock. + +use std::path::Path; + +use crate::compute_logs::_helpers::paths::{ensure_contained, ensure_not_symlink}; +use crate::compute_logs::_helpers::prune::{acquire_lock, release_lock}; +use crate::compute_logs::_helpers::stored::{delete_inactive, inventory}; +use crate::compute_logs::constants::COMPLETE_FILE_NAME; +use crate::compute_logs::errors::ComputeLogError; +use crate::compute_logs::models::{InventoryEntry, PruneOutcome}; + +fn completion_key(entry: &InventoryEntry) -> (i64, &str) { + let instant = entry + .document + .completion + .as_ref() + .map_or(entry.document.started_at.micros, |completion| { + completion.completed_at.micros + }); + (instant, entry.document.invocation_id.as_str()) +} + +/// Delete all but `retain` complete captures, newest kept; incomplete captures always stay. +pub fn prune_captures( + root: &Path, + retain: Option, + is_active: &dyn Fn(&str) -> bool, +) -> Result { + let Some(retain) = retain else { + let captures = inventory(root)?; + let complete = captures.iter().filter(|entry| entry.is_complete).count(); + return Ok(PruneOutcome { + deleted_invocation_ids: Vec::new(), + retained_complete_count: complete, + retained_incomplete_count: captures.len() - complete, + }); + }; + let lock = acquire_lock(root)?; + let result = prune_locked(root, retain, is_active); + release_lock(root, lock); + result +} + +fn prune_locked( + root: &Path, + retain: usize, + is_active: &dyn Fn(&str) -> bool, +) -> Result { + let captures = inventory(root)?; + let incomplete = captures.iter().filter(|entry| !entry.is_complete).count(); + let mut complete: Vec<&InventoryEntry> = + captures.iter().filter(|entry| entry.is_complete).collect(); + complete.sort_by(|left, right| completion_key(right).cmp(&completion_key(left))); + let mut deleted: Vec = Vec::new(); + for entry in complete.iter().skip(retain) { + ensure_not_symlink(&entry.path)?; + ensure_contained(root, &entry.path)?; + let marker = entry.path.join(COMPLETE_FILE_NAME); + ensure_contained(root, &marker)?; + if !marker.is_file() { + continue; + } + let invocation_id = entry.document.invocation_id.as_str(); + delete_inactive(root, invocation_id, is_active(invocation_id))?; + deleted.push(invocation_id.to_owned()); + } + Ok(PruneOutcome { + retained_complete_count: complete.len() - deleted.len(), + deleted_invocation_ids: deleted, + retained_incomplete_count: incomplete, + }) +} diff --git a/crates/sqlbuild-cli/src/compute_logs/main/publish_capture.rs b/crates/sqlbuild-cli/src/compute_logs/main/publish_capture.rs new file mode 100644 index 0000000000..8040ee149d --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/main/publish_capture.rs @@ -0,0 +1,42 @@ +//! `start_capture`: publish one incomplete capture and open its append streams. + +use std::fs::File; +use std::path::Path; + +use crate::compute_logs::_helpers::document::{encode_document, validate_document}; +use crate::compute_logs::_helpers::files::{publish_staged, uuid4_hex}; +use crate::compute_logs::_helpers::paths::{ensure_contained, ensure_not_symlink, find_capture}; +use crate::compute_logs::errors::{ComputeLogError, ComputeLogErrorKind}; +use crate::compute_logs::models::CaptureDocument; + +/// The open stdout, stderr and diagnostics files of the newly published capture. +pub fn publish_capture( + root: &Path, + document: &CaptureDocument, +) -> Result<[File; 3], ComputeLogError> { + validate_document(document)?; + let date_directory = root.join(&document.capture_date); + ensure_not_symlink(&date_directory)?; + std::fs::create_dir_all(&date_directory)?; + ensure_not_symlink(&date_directory)?; + ensure_contained(root, &date_directory)?; + let directory = date_directory.join(&document.invocation_id); + ensure_contained(root, &directory)?; + if find_capture(root, &document.invocation_id)?.is_some() { + return Err(ComputeLogError::rejected( + ComputeLogErrorKind::AlreadyExists, + format!( + "capture already exists for invocation '{}'", + document.invocation_id + ), + )); + } + let staging = date_directory.join(format!(".{}.{}.tmp", document.invocation_id, uuid4_hex()?)); + ensure_contained(root, &staging)?; + std::fs::create_dir(&staging)?; + Ok(publish_staged( + &staging, + &directory, + &encode_document(document), + )?) +} diff --git a/crates/sqlbuild-cli/src/compute_logs/main/read_capture_document.rs b/crates/sqlbuild-cli/src/compute_logs/main/read_capture_document.rs new file mode 100644 index 0000000000..83f64f1526 --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/main/read_capture_document.rs @@ -0,0 +1,30 @@ +//! `get_metadata`: the published metadata of one capture, validated against its directory. + +use std::path::{Path, PathBuf}; + +use crate::compute_logs::_helpers::document::{decode_document, validate_document}; +use crate::compute_logs::_helpers::paths::{ensure_contained, existing_capture}; +use crate::compute_logs::_helpers::stored::validate_directory_document; +use crate::compute_logs::constants::METADATA_FILE_NAME; +use crate::compute_logs::errors::{ComputeLogError, ComputeLogErrorKind}; +use crate::compute_logs::models::CaptureDocument; + +/// The capture's document and directory; read failures become metadata errors. +pub fn read_capture_document( + root: &Path, + invocation_id: &str, +) -> Result<(CaptureDocument, PathBuf), ComputeLogError> { + let directory = existing_capture(root, invocation_id)?; + let path = directory.join(METADATA_FILE_NAME); + ensure_contained(root, &path)?; + let raw = std::fs::read_to_string(&path).map_err(|error| { + ComputeLogError::rejected( + ComputeLogErrorKind::Metadata, + format!("unable to read capture metadata: {error}"), + ) + })?; + let document = decode_document(&raw)?; + validate_document(&document)?; + validate_directory_document(&directory, &document)?; + Ok((document, directory)) +} diff --git a/crates/sqlbuild-cli/src/compute_logs/main/read_stream_chunk.rs b/crates/sqlbuild-cli/src/compute_logs/main/read_stream_chunk.rs new file mode 100644 index 0000000000..b78cda4400 --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/main/read_stream_chunk.rs @@ -0,0 +1,52 @@ +//! `read`: at most `max_bytes` of one stream from an exact byte cursor. + +use std::io::{Read, Seek, SeekFrom}; +use std::path::Path; + +use crate::compute_logs::_helpers::paths::{ + ensure_contained, existing_capture, read_bounds, validate_capture_id, +}; +use crate::compute_logs::constants::COMPLETE_FILE_NAME; +use crate::compute_logs::errors::{ComputeLogError, ComputeLogErrorKind}; +use crate::compute_logs::models::StreamChunk; + +/// The chunk at `cursor` of the capture's `file_name` stream. +pub fn read_stream_chunk( + root: &Path, + (invocation_id, file_name): (&str, &str), + cursor: i64, + max_bytes: i64, +) -> Result { + validate_capture_id(invocation_id)?; + let (cursor, limit) = read_bounds(cursor, max_bytes)?; + let directory = existing_capture(root, invocation_id)?; + let path = directory.join(file_name); + let marker = directory.join(COMPLETE_FILE_NAME); + ensure_contained(root, &path)?; + ensure_contained(root, &marker)?; + let is_complete = marker.is_file(); + if !path.exists() { + return Ok(StreamChunk { + data: Vec::new(), + next_cursor: cursor, + is_complete, + }); + } + let size = std::fs::metadata(&path)?.len(); + if cursor > size { + return Err(ComputeLogError::rejected( + ComputeLogErrorKind::InvalidCursor, + format!("cursor {cursor} is beyond current stream size {size}"), + )); + } + let mut file = std::fs::File::open(&path)?; + file.seek(SeekFrom::Start(cursor))?; + let mut data: Vec = Vec::new(); + Read::by_ref(&mut file).take(limit).read_to_end(&mut data)?; + let next_cursor = file.stream_position()?; + Ok(StreamChunk { + data, + next_cursor, + is_complete, + }) +} diff --git a/crates/sqlbuild-cli/src/compute_logs/main/start_capture.rs b/crates/sqlbuild-cli/src/compute_logs/main/start_capture.rs index 35cf8ab225..c354873ec4 100644 --- a/crates/sqlbuild-cli/src/compute_logs/main/start_capture.rs +++ b/crates/sqlbuild-cli/src/compute_logs/main/start_capture.rs @@ -1,14 +1,10 @@ //! Publish an incomplete capture before the command writes output, as `start_capture` does. -use std::io::{self, Write}; -use std::path::Path; +use std::io; -use crate::compute_logs::_helpers::files::{ - create_new, is_symlink, sync_directory, uuid4_hex, write_metadata, -}; +use crate::compute_logs::_helpers::files::{is_symlink, publish_staged, uuid4_hex}; use crate::compute_logs::_helpers::metadata::metadata_json; use crate::compute_logs::_helpers::timestamps::{epoch_microseconds, utc_date}; -use crate::compute_logs::constants::STREAM_FILE_NAMES; use crate::compute_logs::models::{CaptureStart, OpenCapture}; /// The published capture, or the failure that leaves the command running without one. @@ -33,39 +29,15 @@ pub fn start_capture(start: CaptureStart) -> io::Result { } let staging = date_directory.join(format!(".{invocation_id}.{}.tmp", uuid4_hex()?)); std::fs::create_dir(&staging)?; - match published((&start, &invocation_id), &staging, &directory) { - Ok(streams) => Ok(OpenCapture { - start, - invocation_id, - directory, - streams, - }), - Err(error) => { - let _ = std::fs::remove_dir_all(&staging); - let _ = std::fs::remove_dir_all(&directory); - Err(error) - } - } -} - -fn published( - start: (&CaptureStart, &str), - staging: &Path, - directory: &Path, -) -> io::Result> { - let mut streams: Vec = Vec::new(); - for name in STREAM_FILE_NAMES { - streams.push(create_new(&staging.join(name))?); - } - write_metadata(staging, &metadata_json(start, None))?; - for stream in &mut streams { - stream.flush()?; - stream.sync_all()?; - } - sync_directory(staging)?; - std::fs::rename(staging, directory)?; - if let Some(parent) = directory.parent() { - sync_directory(parent)?; - } - Ok(streams) + let streams = publish_staged( + &staging, + &directory, + &metadata_json((&start, &invocation_id), None), + )?; + Ok(OpenCapture { + start, + invocation_id, + directory, + streams: Vec::from(streams), + }) } diff --git a/crates/sqlbuild-cli/src/compute_logs/main/stream_byte_counts.rs b/crates/sqlbuild-cli/src/compute_logs/main/stream_byte_counts.rs new file mode 100644 index 0000000000..749065f489 --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/main/stream_byte_counts.rs @@ -0,0 +1,21 @@ +//! `get_byte_counts`: the current size of each stream file of one capture. + +use std::path::Path; + +use crate::compute_logs::_helpers::paths::{ensure_contained, existing_capture}; +use crate::compute_logs::constants::STREAM_FILE_NAMES; +use crate::compute_logs::errors::ComputeLogError; + +/// Stdout, stderr and diagnostics sizes; a missing stream file counts as empty. +pub fn stream_byte_counts(root: &Path, invocation_id: &str) -> Result<[u64; 3], ComputeLogError> { + let directory = existing_capture(root, invocation_id)?; + let mut sizes: [u64; 3] = [0; 3]; + for (size, name) in sizes.iter_mut().zip(STREAM_FILE_NAMES) { + let path = directory.join(name); + ensure_contained(root, &path)?; + if path.exists() { + *size = std::fs::metadata(&path)?.len(); + } + } + Ok(sizes) +} diff --git a/crates/sqlbuild-cli/src/compute_logs/main/validate_capture_id.rs b/crates/sqlbuild-cli/src/compute_logs/main/validate_capture_id.rs new file mode 100644 index 0000000000..0f1e333694 --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/main/validate_capture_id.rs @@ -0,0 +1,9 @@ +//! Reject capture identities that are unsafe as local directory names. + +use crate::compute_logs::_helpers::paths; +use crate::compute_logs::errors::ComputeLogError; + +/// Python's `[A-Za-z0-9][A-Za-z0-9_-]{0,127}\Z` capture identity. +pub fn validate_capture_id(invocation_id: &str) -> Result<(), ComputeLogError> { + paths::validate_capture_id(invocation_id) +} diff --git a/crates/sqlbuild-cli/src/compute_logs/mod.rs b/crates/sqlbuild-cli/src/compute_logs/mod.rs index 7366aa64c1..ed4c422fe1 100644 --- a/crates/sqlbuild-cli/src/compute_logs/mod.rs +++ b/crates/sqlbuild-cli/src/compute_logs/mod.rs @@ -2,6 +2,7 @@ pub(crate) mod _helpers; pub mod constants; +pub mod errors; pub mod main; pub mod models; diff --git a/crates/sqlbuild-cli/src/compute_logs/models.rs b/crates/sqlbuild-cli/src/compute_logs/models.rs index 9511254d2e..6ff75166d8 100644 --- a/crates/sqlbuild-cli/src/compute_logs/models.rs +++ b/crates/sqlbuild-cli/src/compute_logs/models.rs @@ -33,3 +33,61 @@ pub struct OpenCapture { /// `stdout.log`, `stderr.log` and `diagnostics.jsonl`. pub streams: Vec, } + +/// One timezone-aware or naive instant of capture metadata. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CaptureInstant { + /// Microseconds since the Unix epoch; meaningful only when `aware`. + pub micros: i64, + pub aware: bool, + /// `isoformat()` with `+00:00` written as `Z`, as `metadata.json` stores it. + pub text: String, +} + +/// What `FinalCaptureMetadata` adds once the capture's streams are closed. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CaptureCompletion { + pub completed_at: CaptureInstant, + pub exit_code: i64, + pub stdout_bytes: i64, + pub stderr_bytes: i64, + pub diagnostics_bytes: i64, +} + +/// Python's `CaptureMetadata`, or `FinalCaptureMetadata` when `completion` is present. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CaptureDocument { + pub format_version: i64, + pub invocation_id: String, + pub command: String, + pub project_dir: String, + pub started_at: CaptureInstant, + pub capture_date: String, + pub target: Option, + pub run_id: Option, + pub completion: Option, +} + +/// One capture beneath a compute log root, as `CaptureInventoryItem`. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct InventoryEntry { + pub document: CaptureDocument, + pub path: PathBuf, + pub is_complete: bool, +} + +/// Raw bytes read from one stream at an exclusive cursor, as `ComputeLogReadChunk`. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct StreamChunk { + pub data: Vec, + pub next_cursor: u64, + pub is_complete: bool, +} + +/// Captures one retention pass deleted and the counts it kept, as `PruneResult`. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct PruneOutcome { + pub deleted_invocation_ids: Vec, + pub retained_complete_count: usize, + pub retained_incomplete_count: usize, +} diff --git a/crates/sqlbuild-cli/src/compute_logs/tests.rs b/crates/sqlbuild-cli/src/compute_logs/tests.rs index 6677eaa86e..215a140006 100644 --- a/crates/sqlbuild-cli/src/compute_logs/tests.rs +++ b/crates/sqlbuild-cli/src/compute_logs/tests.rs @@ -4,5 +4,7 @@ mod capture; mod helpers; #[path = "tests/test_metadata.rs"] mod metadata; +#[path = "tests/test_publication.rs"] +mod publication; #[path = "tests/test_types.rs"] mod test_types; diff --git a/crates/sqlbuild-cli/src/compute_logs/tests/test_metadata.rs b/crates/sqlbuild-cli/src/compute_logs/tests/test_metadata.rs index 94e63b1df9..62a029c48d 100644 --- a/crates/sqlbuild-cli/src/compute_logs/tests/test_metadata.rs +++ b/crates/sqlbuild-cli/src/compute_logs/tests/test_metadata.rs @@ -2,7 +2,7 @@ use std::path::PathBuf; use std::time::{Duration, UNIX_EPOCH}; use crate::compute_logs::_helpers::metadata::metadata_json; -use crate::compute_logs::_helpers::timestamps::{parse_utc_timestamp, utc_timestamp}; +use crate::compute_logs::_helpers::timestamps::{parse_iso_instant, utc_timestamp}; use crate::compute_logs::models::{CaptureEnd, CaptureStart}; use crate::compute_logs::tests::test_types::{MetadataTestCase, TimestampTestCase}; @@ -37,8 +37,14 @@ fn given_utc_instant_when_formatting_then_matches_python_isoformat_and_parses_ba let timestamp = utc_timestamp(test_case.microseconds); assert_eq!( - (timestamp.as_str(), parse_utc_timestamp(×tamp)), - (test_case.expected_timestamp, Some(test_case.microseconds)), + ( + timestamp.as_str(), + parse_iso_instant(×tamp.replace('Z', "+00:00")), + ), + ( + test_case.expected_timestamp, + Some((test_case.microseconds, true)) + ), "{}", test_case.description ); diff --git a/crates/sqlbuild-cli/src/compute_logs/tests/test_publication.rs b/crates/sqlbuild-cli/src/compute_logs/tests/test_publication.rs new file mode 100644 index 0000000000..9104066d63 --- /dev/null +++ b/crates/sqlbuild-cli/src/compute_logs/tests/test_publication.rs @@ -0,0 +1,48 @@ +use crate::compute_logs::_helpers::files::publish_staged; +use crate::compute_logs::tests::test_types::PublicationTestCase; + +#[test] +fn given_failing_staged_step_when_publishing_then_no_capture_directory_remains() { + let test_cases = [ + PublicationTestCase { + description: "a clean staging directory is published", + blocking_entries: &[], + expected_published: true, + expected_entries: 1, + }, + PublicationTestCase { + description: "metadata that cannot replace its destination removes the staging", + blocking_entries: &["metadata.json"], + expected_published: false, + expected_entries: 0, + }, + PublicationTestCase { + description: "a stream file that already exists removes the staging", + blocking_entries: &["stderr.log"], + expected_published: false, + expected_entries: 0, + }, + ]; + + for test_case in test_cases { + let date_directory = tempfile::tempdir().expect("temporary directory"); + let staging = date_directory.path().join(".capture.staging.tmp"); + let directory = date_directory.path().join("capture"); + std::fs::create_dir(&staging).expect("staging directory"); + for name in test_case.blocking_entries { + std::fs::create_dir_all(staging.join(name).join("occupied")).expect("blocking entry"); + } + + let published = publish_staged(&staging, &directory, "{}\n").is_ok(); + let entries = std::fs::read_dir(date_directory.path()) + .map(Iterator::count) + .unwrap_or_default(); + + assert_eq!( + (published, entries), + (test_case.expected_published, test_case.expected_entries), + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-cli/src/compute_logs/tests/test_types.rs b/crates/sqlbuild-cli/src/compute_logs/tests/test_types.rs index 7f41b10a28..6ecc62f743 100644 --- a/crates/sqlbuild-cli/src/compute_logs/tests/test_types.rs +++ b/crates/sqlbuild-cli/src/compute_logs/tests/test_types.rs @@ -20,3 +20,12 @@ pub(crate) struct PruneTestCase { pub(crate) existing: usize, pub(crate) expected_complete: usize, } + +/// A staged capture whose publication fails at one step, and what stays on disk. +pub(crate) struct PublicationTestCase { + pub(crate) description: &'static str, + /// Names created as non-empty directories inside staging before publication. + pub(crate) blocking_entries: &'static [&'static str], + pub(crate) expected_published: bool, + pub(crate) expected_entries: usize, +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/cli/registration.rs b/crates/sqlbuild-python/src/bindings/_helpers/cli/registration.rs index aa0308a836..2857cf1d2b 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/cli/registration.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/cli/registration.rs @@ -7,11 +7,11 @@ use crate::bindings::_helpers::cli::{ progress_text, }; use crate::bindings::_helpers::{ - audit_command, build_command, build_progress, check, clone, command_runtime, contract_command, - cost_command, cost_report, cost_telemetry, dbt_command, dbt_events, dbt_interop, debug_command, - diff_report, format_command, freshness, init_command, janitor, load, plan_command, - presentation_text, query_command, result_output, scenario_command, seed_command, - sql_test_command, workspace_commands, + audit_command, build_command, build_progress, check, clone, command_runtime, compute_logs, + contract_command, cost_command, cost_report, cost_telemetry, dbt_command, dbt_events, + dbt_interop, debug_command, diff_report, format_command, freshness, init_command, janitor, + load, plan_command, presentation_text, query_command, result_output, scenario_command, + seed_command, sql_test_command, workspace_commands, }; pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { @@ -27,6 +27,7 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { plan_command::registration::register(module)?; freshness::registration::register(module)?; clone::registration::register(module)?; + compute_logs::registration::register(module)?; load::registration::register(module)?; result_output::registration::register(module)?; cost_report::registration::register(module)?; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/conversion.rs b/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/conversion.rs new file mode 100644 index 0000000000..5b18e88d17 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/conversion.rs @@ -0,0 +1,252 @@ +//! Compute log metadata models and errors across the Python boundary. + +use pyo3::IntoPyObjectExt; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyErr, PyResult, Python}; +use pyo3::types::{PyInt, PyString, PyStringMethods}; +use sqlbuild_cli::compute_logs::errors::{ComputeLogError, ComputeLogErrorKind}; +use sqlbuild_cli::compute_logs::models::{ + CaptureCompletion, CaptureDocument, CaptureInstant, InventoryEntry, +}; + +use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported}; + +pub(crate) const EXCEPTIONS_MODULE: &str = "sqlbuild.runtime.compute_logs.exceptions"; +pub(crate) const MODELS_MODULE: &str = "sqlbuild.runtime.compute_logs.models"; +pub(crate) const TYPES_MODULE: &str = "sqlbuild.runtime.compute_logs.types"; +const EPOCH_TEXT: &str = "1970-01-01T00:00:00+00:00"; +const UTC_OFFSET: &str = "+00:00"; +const UTC_SUFFIX: &str = "Z"; +const STREAM_VALUES: [&str; 3] = ["stdout", "stderr", "diagnostics"]; + +fn exception_name(kind: ComputeLogErrorKind) -> &'static str { + match kind { + ComputeLogErrorKind::InvalidCaptureId => "InvalidCaptureIdError", + ComputeLogErrorKind::InvalidCursor => "InvalidComputeLogCursorError", + ComputeLogErrorKind::InvalidLimit => "InvalidComputeLogLimitError", + ComputeLogErrorKind::AlreadyExists => "CaptureAlreadyExistsError", + ComputeLogErrorKind::NotFound => "CaptureNotFoundError", + ComputeLogErrorKind::State => "CaptureStateError", + ComputeLogErrorKind::Path => "ComputeLogPathError", + ComputeLogErrorKind::Metadata => "ComputeLogMetadataError", + } +} + +/// The Python exception for a storage rejection of `kind`. +pub(crate) fn rejection(py: Python<'_>, kind: ComputeLogErrorKind, message: &str) -> PyErr { + match imported(py, EXCEPTIONS_MODULE, exception_name(kind)) + .and_then(|class| class.call1((message,))) + { + Ok(error) => PyErr::from_value(error), + Err(error) => error, + } +} + +/// The Python exception a native storage failure raises. +pub(crate) fn storage_error(py: Python<'_>, error: ComputeLogError) -> PyErr { + match error { + ComputeLogError::Rejected(kind, message) => rejection(py, kind, &message), + ComputeLogError::Io(error) => PyErr::from(error), + } +} + +/// Whether `value` is exactly an `int`, as `type(value) is int` checks. +pub(crate) fn is_exact_int(value: &Bound<'_, PyAny>) -> bool { + value.get_type().is(value.py().get_type::()) +} + +/// `invocation_id` as text, rejecting non-strings as invalid capture ids. +pub(crate) fn capture_id(value: &Bound<'_, PyAny>) -> PyResult { + if let Ok(text) = value.downcast::() { + return Ok(text.to_str()?.to_owned()); + } + Err(rejection( + value.py(), + ComputeLogErrorKind::InvalidCaptureId, + "invocation_id must contain only ASCII letters, digits, underscores, or hyphens", + )) +} + +/// The stream index of a `ComputeLogStream` member: stdout, stderr or diagnostics. +pub(crate) fn stream_index(stream: &Bound<'_, PyAny>) -> PyResult { + let py = stream.py(); + if stream.is_instance(&imported(py, TYPES_MODULE, "ComputeLogStream")?)? { + let value: String = stream.getattr("value")?.extract()?; + if let Some(index) = STREAM_VALUES.iter().position(|name| *name == value) { + return Ok(index); + } + } + Err(rejection( + py, + ComputeLogErrorKind::State, + "stream must be a ComputeLogStream", + )) +} + +fn instant(value: &Bound<'_, PyAny>) -> PyResult { + let py = value.py(); + let text: String = value.call_method0("isoformat")?.extract()?; + let text = text.replace(UTC_OFFSET, UTC_SUFFIX); + if value.getattr("tzinfo")?.is_none() || value.call_method0("utcoffset")?.is_none() { + return Ok(CaptureInstant { + micros: 0, + aware: false, + text, + }); + } + let datetime = imported(py, "datetime", "datetime")?; + let epoch = datetime.call_method1("fromisoformat", (EPOCH_TEXT,))?; + let unit = call_with( + &imported(py, "datetime", "timedelta")?, + &[("microseconds", 1_i64.into_bound_py_any(py)?)], + )?; + let micros: i64 = value.sub(epoch)?.floor_div(unit)?.extract()?; + Ok(CaptureInstant { + micros, + aware: true, + text, + }) +} + +fn count(py: Python<'_>, value: &Bound<'_, PyAny>) -> PyResult { + let message = "exit_code and byte counts must be integers"; + if !is_exact_int(value) { + return Err(rejection(py, ComputeLogErrorKind::Metadata, message)); + } + value + .extract() + .map_err(|_| rejection(py, ComputeLogErrorKind::Metadata, message)) +} + +/// A Python `CaptureMetadata` or `FinalCaptureMetadata` as a native capture document. +pub(crate) fn document(metadata: &Bound<'_, PyAny>) -> PyResult { + let py = metadata.py(); + let invocation_id = capture_id(&metadata.getattr("invocation_id")?)?; + let version = metadata.getattr("format_version")?; + let format_version: i64 = if is_exact_int(&version) { + version.extract().unwrap_or(0) + } else { + i64::from(version.eq(1)?) + }; + let completion = + if metadata.is_instance(&imported(py, MODELS_MODULE, "FinalCaptureMetadata")?)? { + Some(CaptureCompletion { + completed_at: instant(&metadata.getattr("completed_at")?)?, + exit_code: count(py, &metadata.getattr("exit_code")?)?, + stdout_bytes: count(py, &metadata.getattr("stdout_bytes")?)?, + stderr_bytes: count(py, &metadata.getattr("stderr_bytes")?)?, + diagnostics_bytes: count(py, &metadata.getattr("diagnostics_bytes")?)?, + }) + } else { + None + }; + Ok(CaptureDocument { + format_version, + invocation_id, + command: metadata.getattr("command")?.extract()?, + project_dir: metadata.getattr("project_dir")?.extract()?, + started_at: instant(&metadata.getattr("started_at")?)?, + capture_date: metadata.getattr("capture_date")?.extract()?, + target: metadata.getattr("target")?.extract()?, + run_id: metadata.getattr("run_id")?.extract()?, + completion, + }) +} + +fn datetime_value<'py>(py: Python<'py>, text: &str) -> PyResult> { + imported(py, "datetime", "datetime")? + .call_method1("fromisoformat", (text.replace(UTC_SUFFIX, UTC_OFFSET),)) + .map_err(|error| { + if error.is_instance_of::(py) { + rejection( + py, + ComputeLogErrorKind::Metadata, + &format!("invalid capture metadata: {}", error.value(py)), + ) + } else { + error + } + }) +} + +/// The Python `CaptureMetadata` or `FinalCaptureMetadata` of a native document. +pub(crate) fn metadata_model<'py>( + py: Python<'py>, + document: &CaptureDocument, +) -> PyResult> { + let mut fields: Vec<(&str, Bound<'py, PyAny>)> = vec![ + ( + "format_version", + document.format_version.into_bound_py_any(py)?, + ), + ( + "invocation_id", + document.invocation_id.as_str().into_bound_py_any(py)?, + ), + ("command", document.command.as_str().into_bound_py_any(py)?), + ( + "project_dir", + document.project_dir.as_str().into_bound_py_any(py)?, + ), + ("started_at", datetime_value(py, &document.started_at.text)?), + ( + "capture_date", + document.capture_date.as_str().into_bound_py_any(py)?, + ), + ("target", document.target.as_deref().into_bound_py_any(py)?), + ("run_id", document.run_id.as_deref().into_bound_py_any(py)?), + ]; + let Some(completion) = &document.completion else { + return call_with(&imported(py, MODELS_MODULE, "CaptureMetadata")?, &fields); + }; + fields.extend([ + ( + "completed_at", + datetime_value(py, &completion.completed_at.text)?, + ), + ("exit_code", completion.exit_code.into_bound_py_any(py)?), + ( + "stdout_bytes", + completion.stdout_bytes.into_bound_py_any(py)?, + ), + ( + "stderr_bytes", + completion.stderr_bytes.into_bound_py_any(py)?, + ), + ( + "diagnostics_bytes", + completion.diagnostics_bytes.into_bound_py_any(py)?, + ), + ]); + call_with( + &imported(py, MODELS_MODULE, "FinalCaptureMetadata")?, + &fields, + ) +} + +/// A Python `CaptureInventoryItem`. +pub(crate) fn inventory_item<'py>( + py: Python<'py>, + entry: &InventoryEntry, +) -> PyResult> { + call_with( + &imported(py, MODELS_MODULE, "CaptureInventoryItem")?, + &[ + ( + "invocation_id", + entry + .document + .invocation_id + .as_str() + .into_bound_py_any(py)?, + ), + ( + "capture_date", + entry.document.capture_date.as_str().into_bound_py_any(py)?, + ), + ("path", entry.path.to_string_lossy().into_bound_py_any(py)?), + ("is_complete", entry.is_complete.into_bound_py_any(py)?), + ("metadata", metadata_model(py, &entry.document)?), + ], + ) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/mod.rs new file mode 100644 index 0000000000..380263eda1 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/mod.rs @@ -0,0 +1,3 @@ +pub(crate) mod conversion; +pub(crate) mod registration; +pub(crate) mod storage; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/registration.rs b/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/registration.rs new file mode 100644 index 0000000000..6ddfa9262d --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/registration.rs @@ -0,0 +1,9 @@ +//! Register the native compute log storage. + +use pyo3::prelude::{Bound, PyModule, PyModuleMethods, PyResult}; + +use crate::bindings::_helpers::compute_logs::storage::LocalFilesystemComputeLogStorage; + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_class::() +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/storage.rs b/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/storage.rs new file mode 100644 index 0000000000..c33ecdee32 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/storage.rs @@ -0,0 +1,431 @@ +//! `LocalFilesystemComputeLogStorage`: the project-local compute log backend. + +use std::collections::{HashMap, HashSet}; +use std::fs::File; +use std::io::{Seek, Write}; +use std::path::PathBuf; +use std::sync::{Mutex, MutexGuard}; + +use pyo3::IntoPyObjectExt; +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::{Bound, Py, PyAny, PyAnyMethods, PyRef, PyResult, Python, pyclass, pymethods}; +use pyo3::types::{PyBytes, PyBytesMethods, PyDict, PyString, PyStringMethods, PyTuple}; +use sqlbuild_cli::compute_logs::constants::{ + DEFAULT_RETENTION_COUNT, MAX_READ_BYTES, STREAM_FILE_NAMES, +}; +use sqlbuild_cli::compute_logs::errors::{ComputeLogError, ComputeLogErrorKind}; +use sqlbuild_cli::compute_logs::main::capture_inventory::capture_inventory; +use sqlbuild_cli::compute_logs::main::capture_is_complete::capture_is_complete; +use sqlbuild_cli::compute_logs::main::check_completion::check_completion; +use sqlbuild_cli::compute_logs::main::complete_capture::complete_capture; +use sqlbuild_cli::compute_logs::main::delete_capture::delete_capture; +use sqlbuild_cli::compute_logs::main::open_log_root::open_log_root; +use sqlbuild_cli::compute_logs::main::prune_captures::prune_captures; +use sqlbuild_cli::compute_logs::main::publish_capture::publish_capture; +use sqlbuild_cli::compute_logs::main::read_capture_document::read_capture_document; +use sqlbuild_cli::compute_logs::main::read_stream_chunk::read_stream_chunk; +use sqlbuild_cli::compute_logs::main::stream_byte_counts::stream_byte_counts; +use sqlbuild_cli::compute_logs::main::validate_capture_id::validate_capture_id; + +use crate::bindings::_helpers::compile_driver::python_calls::{ + KeywordArguments, call_with, imported, +}; +use crate::bindings::_helpers::compute_logs::conversion::{ + MODELS_MODULE, capture_id, document, inventory_item, is_exact_int, metadata_model, rejection, + storage_error, stream_index, +}; + +const DIAGNOSTIC_MODEL_MODULE: &str = "sqlbuild.runtime.observability.models"; +const DIAGNOSTIC_JSON_MODULE: &str = "sqlbuild.runtime.observability.main.diagnostic_log_to_json"; +const DIAGNOSTICS_STREAM: usize = 2; + +/// Open stream files of the captures this instance started, and whether it is closed. +#[derive(Default)] +struct StorageState { + writers: HashMap, + closed: bool, +} + +/// Ordered fsync capture storage subject to underlying filesystem and hardware guarantees. +#[pyclass(module = "sqlbuild._native", frozen, subclass)] +pub(crate) struct LocalFilesystemComputeLogStorage { + root: PathBuf, + retention_count: Option, + state: Mutex, +} + +/// The stdout, stderr and diagnostics byte counts of open streams, as `tell()` reports them. +fn stream_positions(streams: &mut [File; 3]) -> PyResult<[i64; 3]> { + let mut counts: [i64; 3] = [0; 3]; + for (count, stream) in counts.iter_mut().zip(streams.iter_mut()) { + *count = i64::try_from(stream.stream_position()?).unwrap_or(i64::MAX); + } + Ok(counts) +} + +fn not_active(py: Python<'_>) -> pyo3::PyErr { + rejection( + py, + ComputeLogErrorKind::State, + "capture is not active in this storage instance", + ) +} + +fn checked(py: Python<'_>, result: Result) -> PyResult { + result.map_err(|error| storage_error(py, error)) +} + +/// `value` as an integer when it is exactly an `int` that fits, else `fallback`. +fn exact_count(value: &Bound<'_, PyAny>, fallback: i64) -> i64 { + if !is_exact_int(value) { + return fallback; + } + value.extract().unwrap_or(fallback) +} + +impl LocalFilesystemComputeLogStorage { + fn locked(&self) -> PyResult> { + self.state + .lock() + .map_err(|_| PyRuntimeError::new_err("compute log storage lock is poisoned")) + } + + fn ensure_open(&self, py: Python<'_>) -> PyResult<()> { + if self.locked()?.closed { + return Err(rejection( + py, + ComputeLogErrorKind::State, + "compute log storage is closed", + )); + } + Ok(()) + } + + fn append_bytes( + &self, + py: Python<'_>, + invocation_id: &str, + stream: usize, + data: &[u8], + ) -> PyResult<()> { + let mut state = self.locked()?; + let Some(writers) = state.writers.get_mut(invocation_id) else { + return Err(not_active(py)); + }; + writers[stream].write_all(data)?; + Ok(()) + } + + fn byte_counts(&self, py: Python<'_>, invocation_id: &Bound<'_, PyAny>) -> PyResult<[u64; 3]> { + let invocation_id = capture_id(invocation_id)?; + checked(py, stream_byte_counts(&self.root, &invocation_id)) + } +} + +#[pymethods] +impl LocalFilesystemComputeLogStorage { + #[new] + #[pyo3(signature = (*, project_dir, root = None, retention_count = Some(DEFAULT_RETENTION_COUNT as i64)))] + fn new( + py: Python<'_>, + project_dir: PathBuf, + root: Option, + retention_count: Option, + ) -> PyResult { + let root = checked( + py, + open_log_root(&project_dir, root.as_deref(), retention_count), + )?; + Ok(Self { + root, + retention_count, + state: Mutex::new(StorageState::default()), + }) + } + + /// Return the resolved configured log root. + #[getter] + fn root<'py>(&self, py: Python<'py>) -> PyResult> { + imported(py, "pathlib", "Path")?.call1((self.root.as_path(),)) + } + + /// Create one incomplete capture and open its serialized append streams. + fn start_capture(&self, py: Python<'_>, metadata: &Bound<'_, PyAny>) -> PyResult<()> { + self.ensure_open(py)?; + let document = document(metadata)?; + let mut state = self.locked()?; + let writers = checked(py, publish_capture(&self.root, &document))?; + state.writers.insert(document.invocation_id, writers); + Ok(()) + } + + /// Append exact raw bytes and make them visible to concurrent readers. + #[pyo3(signature = (*, invocation_id, stream, data))] + fn append( + &self, + py: Python<'_>, + invocation_id: &Bound<'_, PyAny>, + stream: &Bound<'_, PyAny>, + data: &Bound<'_, PyAny>, + ) -> PyResult<()> { + self.ensure_open(py)?; + let invocation_id = capture_id(invocation_id)?; + checked(py, validate_capture_id(&invocation_id))?; + let stream = stream_index(stream)?; + let Ok(bytes) = data.downcast::() else { + return Err(rejection( + py, + ComputeLogErrorKind::State, + "compute log append data must be bytes", + )); + }; + self.append_bytes(py, &invocation_id, stream, bytes.as_bytes()) + } + + /// Append one validated deterministic UTF-8 diagnostic JSON line. + #[pyo3(signature = (*, invocation_id, record))] + fn append_diagnostic( + &self, + py: Python<'_>, + invocation_id: &Bound<'_, PyAny>, + record: &Bound<'_, PyAny>, + ) -> PyResult<()> { + if !record.is_instance(&imported(py, DIAGNOSTIC_MODEL_MODULE, "DiagnosticLog")?)? { + return Err(rejection( + py, + ComputeLogErrorKind::State, + "diagnostic append requires DiagnosticLog", + )); + } + let encoded: Bound<'_, PyAny> = + imported(py, DIAGNOSTIC_JSON_MODULE, "diagnostic_log_to_json")?.call1((record,))?; + let mut line: Vec = encoded + .downcast::()? + .to_str()? + .as_bytes() + .to_vec(); + line.push(b'\n'); + self.ensure_open(py)?; + let invocation_id = capture_id(invocation_id)?; + checked(py, validate_capture_id(&invocation_id))?; + self.append_bytes(py, &invocation_id, DIAGNOSTICS_STREAM, &line) + } + + /// Read at most the public byte limit from an exact stream offset. + #[pyo3(signature = (**options))] + fn read<'py>( + &self, + py: Python<'py>, + options: Option>, + ) -> PyResult> { + let options = KeywordArguments::read( + py, + "read", + options, + (&["invocation_id", "stream"], &["cursor", "max_bytes"]), + )?; + let invocation_id = capture_id(&options.value("invocation_id")?)?; + checked(py, validate_capture_id(&invocation_id))?; + let stream = stream_index(&options.value("stream")?)?; + let cursor = match options.contains("cursor")? { + true => exact_count(&options.value("cursor")?, -1), + false => 0, + }; + let max_bytes = match options.contains("max_bytes")? { + true => exact_count(&options.value("max_bytes")?, 0), + false => i64::try_from(MAX_READ_BYTES).unwrap_or(i64::MAX), + }; + let chunk = checked( + py, + read_stream_chunk( + &self.root, + (&invocation_id, STREAM_FILE_NAMES[stream]), + cursor, + max_bytes, + ), + )?; + call_with( + &imported(py, MODELS_MODULE, "ComputeLogReadChunk")?, + &[ + ("data", PyBytes::new(py, &chunk.data).into_any()), + ("next_cursor", chunk.next_cursor.into_bound_py_any(py)?), + ("is_complete", chunk.is_complete.into_bound_py_any(py)?), + ], + ) + } + + /// Read the currently atomically published metadata document. + #[pyo3(name = "get_metadata", signature = (*, invocation_id))] + fn metadata<'py>( + &self, + py: Python<'py>, + invocation_id: &Bound<'py, PyAny>, + ) -> PyResult> { + let invocation_id = capture_id(invocation_id)?; + let (document, _) = checked(py, read_capture_document(&self.root, &invocation_id))?; + metadata_model(py, &document) + } + + /// Return the current raw byte count for one stream. + #[pyo3(name = "get_byte_count", signature = (*, invocation_id, stream))] + fn byte_count( + &self, + py: Python<'_>, + invocation_id: &Bound<'_, PyAny>, + stream: &Bound<'_, PyAny>, + ) -> PyResult { + let stream = stream_index(stream)?; + Ok(self.byte_counts(py, invocation_id)?[stream]) + } + + /// Return one lock-consistent snapshot of all stream sizes. + #[pyo3(name = "get_byte_counts", signature = (*, invocation_id))] + fn all_byte_counts<'py>( + &self, + py: Python<'py>, + invocation_id: &Bound<'py, PyAny>, + ) -> PyResult> { + let [stdout_bytes, stderr_bytes, diagnostics_bytes] = + self.byte_counts(py, invocation_id)?; + call_with( + &imported(py, MODELS_MODULE, "CaptureByteCounts")?, + &[ + ("stdout_bytes", stdout_bytes.into_bound_py_any(py)?), + ("stderr_bytes", stderr_bytes.into_bound_py_any(py)?), + ( + "diagnostics_bytes", + diagnostics_bytes.into_bound_py_any(py)?, + ), + ], + ) + } + + /// Return capture-close evidence without asserting execution lifecycle status. + #[pyo3(name = "is_complete", signature = (*, invocation_id))] + fn complete_marker(&self, py: Python<'_>, invocation_id: &Bound<'_, PyAny>) -> PyResult { + let invocation_id = capture_id(invocation_id)?; + checked(py, capture_is_complete(&self.root, &invocation_id)) + } + + /// Flush, fsync, and close streams before publishing final metadata and marker. + fn mark_complete(&self, py: Python<'_>, metadata: &Bound<'_, PyAny>) -> PyResult<()> { + self.ensure_open(py)?; + let document = document(metadata)?; + let mut state = self.locked()?; + let counts = match state.writers.get_mut(&document.invocation_id) { + Some(streams) => Some(stream_positions(streams)?), + None => None, + }; + let directory = checked(py, check_completion(&self.root, &document, counts))?; + let Some(streams) = state.writers.get(&document.invocation_id) else { + return Err(not_active(py)); + }; + for stream in streams { + stream.sync_all()?; + } + let Some(streams) = state.writers.remove(&document.invocation_id) else { + return Err(not_active(py)); + }; + drop(state); + checked(py, complete_capture(&directory, &document, streams)) + } + + /// Delete a complete inactive capture without following symlinks. + #[pyo3(signature = (*, invocation_id))] + fn delete(&self, py: Python<'_>, invocation_id: &Bound<'_, PyAny>) -> PyResult<()> { + let invocation_id = capture_id(invocation_id)?; + let state = self.locked()?; + let active = state.writers.contains_key(&invocation_id); + let result = delete_capture(&self.root, &invocation_id, active); + drop(state); + checked(py, result) + } + + /// Retain the newest complete captures while never deleting incomplete captures. + #[pyo3(signature = (retain_count = None))] + fn prune<'py>( + &self, + py: Python<'py>, + retain_count: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + let retain: Option = match retain_count { + None => self + .retention_count + .map(|count| usize::try_from(count).unwrap_or(0)), + Some(value) => Some(usize::try_from(exact_count(value, -1)).map_err(|_| { + rejection( + py, + ComputeLogErrorKind::InvalidLimit, + "retain_count must be nonnegative or None", + ) + })?), + }; + let active: HashSet = self.locked()?.writers.keys().cloned().collect(); + let root = self.root.clone(); + let outcome = checked( + py, + py.detach(|| prune_captures(&root, retain, &|id: &str| active.contains(id))), + )?; + call_with( + &imported(py, MODELS_MODULE, "PruneResult")?, + &[ + ( + "deleted_invocation_ids", + PyTuple::new(py, &outcome.deleted_invocation_ids)?.into_any(), + ), + ( + "retained_complete_count", + outcome.retained_complete_count.into_bound_py_any(py)?, + ), + ( + "retained_incomplete_count", + outcome.retained_incomplete_count.into_bound_py_any(py)?, + ), + ], + ) + } + + /// Inventory complete and abandoned captures without synthesizing lifecycle state. + fn inventory<'py>(&self, py: Python<'py>) -> PyResult> { + let entries = checked(py, capture_inventory(&self.root))?; + let mut items: Vec> = Vec::new(); + for entry in &entries { + items.push(inventory_item(py, entry)?); + } + let complete_count = entries.iter().filter(|entry| entry.is_complete).count(); + call_with( + &imported(py, MODELS_MODULE, "CaptureInventory")?, + &[ + ("captures", PyTuple::new(py, items)?.into_any()), + ("complete_count", complete_count.into_bound_py_any(py)?), + ( + "incomplete_count", + (entries.len() - complete_count).into_bound_py_any(py)?, + ), + ], + ) + } + + /// Close active writers, deliberately leaving their captures incomplete. + fn close(&self) -> PyResult<()> { + let mut state = self.locked()?; + state.closed = true; + state.writers.clear(); + Ok(()) + } + + /// Dispose process-local resources without deleting retained logs. + fn dispose(&self) -> PyResult<()> { + self.close() + } + + fn __enter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> { + slf + } + + #[pyo3(signature = (*_exc_info))] + fn __exit__(&self, _exc_info: &Bound<'_, PyTuple>) -> PyResult>> { + self.close()?; + Ok(None) + } +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/mod.rs index a63b0b867b..290dead3f9 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/mod.rs @@ -17,6 +17,7 @@ pub(crate) mod compile_command; pub(crate) mod compile_driver; pub(crate) mod compile_outputs; pub(crate) mod compiled_project; +pub(crate) mod compute_logs; pub(crate) mod contract_command; pub(crate) mod contracts; pub(crate) mod cost_command; diff --git a/python-boundary.toml b/python-boundary.toml index b803e03e59..a80f8da9de 100644 --- a/python-boundary.toml +++ b/python-boundary.toml @@ -235,9 +235,44 @@ reason = "executor" note = "Deadline the executor scheduler and warehouse cursors enforce." [[stays]] -path = "runtime/compute_logs" -reason = "executor" -note = "Captures compute logs while the executor runs models, loaders and tasks." +path = "runtime/compute_logs/classes" +reason = "user-code" +note = "Rule 1 capture shim: sys.stdout/sys.stderr tees and a logging handler that record Python output from user code, adapters and the executor (_dispatch_with_compute_logs.py, compile_command/reuse_attempt.rs); storage itself is native LocalFilesystemComputeLogStorage." + +[[stays]] +path = "runtime/compute_logs/main" +reason = "user-code" +note = "Rule 1 capture shim entry: installs the stream tees around the CLI command (_dispatch_with_compute_logs.py)." + +[[stays]] +path = "runtime/compute_logs/_helpers" +reason = "user-code" +note = "Rule 1 tee shim: how much of a write the wrapped Python stream accepted (text_tee.py, binary_tee.py)." + +[[stays]] +path = "runtime/compute_logs/__init__.py" +reason = "public-api" +note = "Package marker of the public compute log shapes below." + +[[stays]] +path = "runtime/compute_logs/models.py" +reason = "public-api" +note = "Rule 3 shapes: capture metadata and result dataclasses re-exported by sqlbuild.compute_logs." + +[[stays]] +path = "runtime/compute_logs/types.py" +reason = "public-api" +note = "Rule 3 shapes: ComputeLogStream and the ComputeLogStorage protocol re-exported by sqlbuild.compute_logs." + +[[stays]] +path = "runtime/compute_logs/exceptions.py" +reason = "public-api" +note = "Rule 3 shapes: public storage errors re-exported by sqlbuild.compute_logs and raised by the native storage." + +[[stays]] +path = "runtime/compute_logs/constants.py" +reason = "public-api" +note = "Rule 3 shapes: format version, retention and read-limit constants re-exported by sqlbuild.compute_logs." [[stays]] path = "runtime/event_exporting" diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index 792d2d09dc..46672425f5 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -11,7 +11,8 @@ from decimal import Decimal from enum import StrEnum from pathlib import Path, PurePath from threading import Event -from typing import Any, ClassVar, Literal, Protocol, TextIO, TypedDict +from types import TracebackType +from typing import Any, ClassVar, Literal, Protocol, Self, TextIO, TypedDict from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.adapter.contract.classes.strict_adapter import StrictAdapter @@ -341,10 +342,19 @@ from sqlbuild.rule_engine.models import ( SqlExpansionReuse, SqlRulesEvaluation, ) +from sqlbuild.runtime.compute_logs.models import ( + CaptureByteCounts, + CaptureInventory, + CaptureMetadata, + ComputeLogReadChunk, + FinalCaptureMetadata, + PruneResult, +) +from sqlbuild.runtime.compute_logs.types import ComputeLogStream from sqlbuild.runtime.contracts.models import ConnectionHooks from sqlbuild.runtime.contracts.types import ExecutionResourceKind from sqlbuild.runtime.observability.classes.event_dispatcher import EventDispatcher -from sqlbuild.runtime.observability.models import LifecycleEvent +from sqlbuild.runtime.observability.models import DiagnosticLog, LifecycleEvent from sqlbuild.spec.contracts.models import ( CostConfig, DbtConfig, @@ -3441,6 +3451,48 @@ class CliObservabilityScope: def __exit__(self, *exc_info: object) -> bool: ... def cli_observability_scope(*, args: CliNamespace, project_dir: Path) -> CliObservabilityScope: ... + +class LocalFilesystemComputeLogStorage: + """Ordered fsync capture storage subject to underlying filesystem and hardware guarantees.""" + + def __init__( + self, + *, + project_dir: Path, + root: Path | None = None, + retention_count: int | None = 20, + ) -> None: ... + @property + def root(self) -> Path: ... + def start_capture(self, metadata: CaptureMetadata) -> None: ... + def append(self, *, invocation_id: str, stream: ComputeLogStream, data: bytes) -> None: ... + def append_diagnostic(self, *, invocation_id: str, record: DiagnosticLog) -> None: ... + def read( + self, + *, + invocation_id: str, + stream: ComputeLogStream, + cursor: int = 0, + max_bytes: int = 1_048_576, + ) -> ComputeLogReadChunk: ... + def get_metadata(self, *, invocation_id: str) -> CaptureMetadata | FinalCaptureMetadata: ... + def get_byte_count(self, *, invocation_id: str, stream: ComputeLogStream) -> int: ... + def get_byte_counts(self, *, invocation_id: str) -> CaptureByteCounts: ... + def is_complete(self, *, invocation_id: str) -> bool: ... + def mark_complete(self, metadata: FinalCaptureMetadata) -> None: ... + def delete(self, *, invocation_id: str) -> None: ... + def prune(self, retain_count: int | None = None) -> PruneResult: ... + def inventory(self) -> CaptureInventory: ... + def close(self) -> None: ... + def dispose(self) -> None: ... + def __enter__(self) -> Self: ... + def __exit__( + self, + exception_type: type[BaseException] | None, + exception: BaseException | None, + traceback: TracebackType | None, + ) -> bool | None: ... + def current_native_progress_projector() -> NativeProgressProjector | None: ... class ConnectionProgressReporter: diff --git a/src/sqlbuild/cli/commands/main/entrypoint/_dispatch_with_compute_logs.py b/src/sqlbuild/cli/commands/main/entrypoint/_dispatch_with_compute_logs.py index 93f9a1c0ad..a8cc32b3bd 100644 --- a/src/sqlbuild/cli/commands/main/entrypoint/_dispatch_with_compute_logs.py +++ b/src/sqlbuild/cli/commands/main/entrypoint/_dispatch_with_compute_logs.py @@ -6,20 +6,17 @@ from datetime import UTC, datetime from pathlib import Path -from sqlbuild._native import CliNamespace, supports_color +from sqlbuild._native import CliNamespace, LocalFilesystemComputeLogStorage, supports_color from sqlbuild.cli.commands._helpers.entry.compute_log_diagnostics import ( log_compute_capture_failure, ) from sqlbuild.diagnostics.classes.invocation_diagnostic_routing import InvocationDiagnosticRouting from sqlbuild.diagnostics.main.configure import configure_diagnostics from sqlbuild.diagnostics.models import DiagnosticRoutingOptions -from sqlbuild.runtime.compute_logs.classes.local_filesystem_compute_log_storage import ( - LocalFilesystemComputeLogStorage, -) -from sqlbuild.runtime.compute_logs.classes.scoped_compute_log_capture import ( - ScopedComputeLogCapture, -) from sqlbuild.runtime.compute_logs.constants import COMPUTE_LOG_FORMAT_VERSION +from sqlbuild.runtime.compute_logs.main.run_with_compute_log_capture import ( + run_with_compute_log_capture, +) from sqlbuild.runtime.compute_logs.models import CaptureMetadata from sqlbuild.runtime.observability.models import ExecutionIdentity @@ -72,12 +69,12 @@ def dispatch_with_compute_logs( for failure, channel in failures: log_compute_capture_failure(error=failure, channel=channel) return operation() - capture: ScopedComputeLogCapture = ScopedComputeLogCapture( + return run_with_compute_log_capture( storage=storage, metadata=metadata, failure_callback=lambda error, channel: log_compute_capture_failure( error=error, channel=channel ), routing_options=routing_options, + operation=operation, ) - return capture.run(operation=operation) diff --git a/src/sqlbuild/compute_logs.py b/src/sqlbuild/compute_logs.py index 353214b936..f9581b03b2 100644 --- a/src/sqlbuild/compute_logs.py +++ b/src/sqlbuild/compute_logs.py @@ -2,9 +2,7 @@ from pathlib import Path -from sqlbuild.runtime.compute_logs.classes.local_filesystem_compute_log_storage import ( - LocalFilesystemComputeLogStorage, -) +from sqlbuild._native import LocalFilesystemComputeLogStorage from sqlbuild.runtime.compute_logs.constants import ( COMPUTE_LOG_FORMAT_VERSION, DEFAULT_RETENTION_COUNT, @@ -21,9 +19,6 @@ InvalidComputeLogCursorError, InvalidComputeLogLimitError, ) -from sqlbuild.runtime.compute_logs.main.open_local_compute_log_storage import ( - open_local_compute_log_storage as _open_local_compute_log_storage, -) from sqlbuild.runtime.compute_logs.models import ( CaptureByteCounts, CaptureInventory, @@ -73,7 +68,7 @@ def open_local_compute_log_storage( ) -> LocalFilesystemComputeLogStorage: """Open the supported project-local backend, defaulting to project logs.""" - return _open_local_compute_log_storage( + return LocalFilesystemComputeLogStorage( project_dir=project_dir, root=root, retention_count=retention_count, diff --git a/src/sqlbuild/runtime/compute_logs/README.md b/src/sqlbuild/runtime/compute_logs/README.md index bc8c8b8e69..853d0d2f79 100644 --- a/src/sqlbuild/runtime/compute_logs/README.md +++ b/src/sqlbuild/runtime/compute_logs/README.md @@ -1,7 +1,8 @@ # Compute log storage contract -This package owns backend-neutral raw invocation stream storage and the host-local filesystem -implementation. It stores exact stdout/stderr bytes and structured `DiagnosticLog` JSONL separately +This package owns the backend-neutral compute log shapes and the Python stream capture around a +command. The host-local filesystem storage is native (`sqlbuild._native.LocalFilesystemComputeLogStorage`, +built on `crates/sqlbuild-cli/src/compute_logs`) and is re-exported by `sqlbuild.compute_logs`. It stores exact stdout/stderr bytes and structured `DiagnosticLog` JSONL separately from canonical lifecycle history. The default root is `logs/`, with captures at diff --git a/src/sqlbuild/runtime/compute_logs/_helpers/metadata.py b/src/sqlbuild/runtime/compute_logs/_helpers/metadata.py deleted file mode 100644 index 99c0029977..0000000000 --- a/src/sqlbuild/runtime/compute_logs/_helpers/metadata.py +++ /dev/null @@ -1,76 +0,0 @@ -"""Compute log metadata JSON encoding and validation.""" - -from __future__ import annotations - -import json -from datetime import datetime -from typing import Any - -from sqlbuild.runtime.compute_logs.constants import COMPLETED_AT_FIELD, COMPUTE_LOG_FORMAT_VERSION -from sqlbuild.runtime.compute_logs.exceptions import ComputeLogMetadataError -from sqlbuild.runtime.compute_logs.models import CaptureMetadata, FinalCaptureMetadata - - -def metadata_to_json(metadata: CaptureMetadata | FinalCaptureMetadata) -> str: - """Serialize capture metadata deterministically without raw command arguments.""" - - payload: dict[str, object] = { - "format_version": metadata.format_version, - "invocation_id": metadata.invocation_id, - "command": metadata.command, - "project_dir": metadata.project_dir, - "started_at": _timestamp(metadata.started_at), - "capture_date": metadata.capture_date, - "target": metadata.target, - "run_id": metadata.run_id, - } - if isinstance(metadata, FinalCaptureMetadata): - payload.update( - { - "completed_at": _timestamp(metadata.completed_at), - "exit_code": metadata.exit_code, - "stdout_bytes": metadata.stdout_bytes, - "stderr_bytes": metadata.stderr_bytes, - "diagnostics_bytes": metadata.diagnostics_bytes, - } - ) - return json.dumps(payload, sort_keys=True, separators=(",", ":"), ensure_ascii=False) + "\n" - - -def metadata_from_json(raw_json: str) -> CaptureMetadata | FinalCaptureMetadata: - """Decode and validate initial or final capture metadata.""" - - try: - payload: Any = json.loads(raw_json) - if not isinstance(payload, dict): - raise ComputeLogMetadataError("capture metadata must be a JSON object") - common: dict[str, object] = { - "format_version": payload["format_version"], - "invocation_id": payload["invocation_id"], - "command": payload["command"], - "project_dir": payload["project_dir"], - "started_at": datetime.fromisoformat(payload["started_at"].replace("Z", "+00:00")), - "capture_date": payload["capture_date"], - "target": payload.get("target"), - "run_id": payload.get("run_id"), - } - if common["format_version"] != COMPUTE_LOG_FORMAT_VERSION: - raise ComputeLogMetadataError("unsupported compute log format version") - if COMPLETED_AT_FIELD not in payload: - return CaptureMetadata(**common) # type: ignore[arg-type] - return FinalCaptureMetadata( - **common, # type: ignore[arg-type] - completed_at=datetime.fromisoformat(payload["completed_at"].replace("Z", "+00:00")), - exit_code=payload["exit_code"], - stdout_bytes=payload["stdout_bytes"], - stderr_bytes=payload["stderr_bytes"], - diagnostics_bytes=payload["diagnostics_bytes"], - ) - except ComputeLogMetadataError: - raise - except (AttributeError, KeyError, TypeError, ValueError, json.JSONDecodeError) as error: - raise ComputeLogMetadataError(f"invalid capture metadata: {error}") from error - - -def _timestamp(value: datetime) -> str: - return value.isoformat().replace("+00:00", "Z") diff --git a/src/sqlbuild/runtime/compute_logs/classes/local_filesystem_compute_log_storage.py b/src/sqlbuild/runtime/compute_logs/classes/local_filesystem_compute_log_storage.py deleted file mode 100644 index 0d358daf17..0000000000 --- a/src/sqlbuild/runtime/compute_logs/classes/local_filesystem_compute_log_storage.py +++ /dev/null @@ -1,616 +0,0 @@ -"""Project-local filesystem compute log storage.""" - -from __future__ import annotations - -import os -import re -import shutil -import threading -from datetime import UTC -from pathlib import Path -from types import TracebackType -from typing import BinaryIO, Literal, Self, overload -from uuid import uuid4 - -from filelock import FileLock - -from sqlbuild.runtime.compute_logs._helpers.metadata import metadata_from_json, metadata_to_json -from sqlbuild.runtime.compute_logs.constants import ( - COMPLETE_FILE_NAME, - COMPUTE_LOG_FORMAT_VERSION, - DEFAULT_RETENTION_COUNT, - MAX_READ_BYTES, - METADATA_FILE_NAME, - PRUNE_LOCK_FILE_NAME, -) -from sqlbuild.runtime.compute_logs.exceptions import ( - CaptureAlreadyExistsError, - CaptureNotFoundError, - CaptureStateError, - ComputeLogMetadataError, - ComputeLogPathError, - InvalidCaptureIdError, - InvalidComputeLogCursorError, - InvalidComputeLogLimitError, -) -from sqlbuild.runtime.compute_logs.models import ( - CaptureByteCounts, - CaptureInventory, - CaptureInventoryItem, - CaptureMetadata, - ComputeLogReadChunk, - FinalCaptureMetadata, - PruneResult, -) -from sqlbuild.runtime.compute_logs.types import ByteCursor, ComputeLogStream -from sqlbuild.runtime.observability.main.diagnostic_log_to_json import diagnostic_log_to_json -from sqlbuild.runtime.observability.models import DiagnosticLog - -_CAPTURE_ID_PATTERN: re.Pattern[str] = re.compile(r"[A-Za-z0-9][A-Za-z0-9_-]{0,127}\Z") -_DATE_PATTERN: re.Pattern[str] = re.compile(r"\d{4}-\d{2}-\d{2}\Z") -_STREAM_FILES: dict[ComputeLogStream, str] = { - ComputeLogStream.STDOUT: "stdout.log", - ComputeLogStream.STDERR: "stderr.log", - ComputeLogStream.DIAGNOSTICS: "diagnostics.jsonl", -} - - -class LocalFilesystemComputeLogStorage: - """Ordered fsync capture storage subject to underlying filesystem and hardware guarantees.""" - - def __init__( - self, - *, - project_dir: Path, - root: Path | None = None, - retention_count: int | None = DEFAULT_RETENTION_COUNT, - ) -> None: - if retention_count is not None and retention_count < 0: - raise InvalidComputeLogLimitError("retention_count must be nonnegative or None") - configured_root: Path = project_dir / "logs" if root is None else root - if configured_root.is_symlink(): - raise ComputeLogPathError("configured compute log root cannot be a symlink") - configured_root.mkdir(parents=True, exist_ok=True) - if configured_root.is_symlink(): - raise ComputeLogPathError("configured compute log root cannot be a symlink") - self._root: Path = configured_root.resolve(strict=True) - if not self._root.is_dir(): - raise ComputeLogPathError("compute log root must be a directory") - self._retention_count: int | None = retention_count - self._writers: dict[str, dict[ComputeLogStream, BinaryIO]] = {} - self._locks: dict[str, threading.RLock] = {} - self._state_lock: threading.RLock = threading.RLock() - self._closed: bool = False - - @property - def root(self) -> Path: - """Return the resolved configured log root.""" - - return self._root - - def start_capture(self, metadata: CaptureMetadata) -> None: - """Create one incomplete capture and open its serialized append streams.""" - - self._ensure_open() - self._validate_metadata(metadata=metadata) - date_dir: Path = self._root / metadata.capture_date - self._ensure_not_symlink(path=date_dir) - date_dir.mkdir(parents=True, exist_ok=True) - self._ensure_not_symlink(path=date_dir) - self._ensure_contained(path=date_dir) - capture_dir: Path = date_dir / metadata.invocation_id - self._ensure_contained(path=capture_dir) - with self._state_lock: - existing: Path | None = self._find_capture( - invocation_id=metadata.invocation_id, missing_ok=True - ) - if existing is not None: - raise CaptureAlreadyExistsError( - f"capture already exists for invocation {metadata.invocation_id!r}" - ) - staging_dir: Path = date_dir / f".{metadata.invocation_id}.{uuid4().hex}.tmp" - self._ensure_contained(path=staging_dir) - staging_dir.mkdir(parents=False, exist_ok=False) - writers: dict[ComputeLogStream, BinaryIO] = {} - published: bool = False - try: - for stream, file_name in _STREAM_FILES.items(): - writers[stream] = (staging_dir / file_name).open("xb") - self._write_metadata(capture_dir=staging_dir, metadata=metadata) - for writer in writers.values(): - writer.flush() - os.fsync(writer.fileno()) - self._fsync_directory(staging_dir) - staging_dir.rename(capture_dir) - published = True - self._fsync_directory(date_dir) - except BaseException: - for writer in writers.values(): - try: - if not writer.closed: - writer.close() - except Exception as cleanup_error: - _ = cleanup_error - failed_dir: Path = capture_dir if published else staging_dir - shutil.rmtree(failed_dir, ignore_errors=True) - raise - self._writers[metadata.invocation_id] = writers - self._locks[metadata.invocation_id] = threading.RLock() - - def append(self, *, invocation_id: str, stream: ComputeLogStream, data: bytes) -> None: - """Append exact raw bytes and make them visible to concurrent readers.""" - - self._ensure_open() - self._validate_capture_id(invocation_id) - if not isinstance(stream, ComputeLogStream): - raise CaptureStateError("stream must be a ComputeLogStream") - if not isinstance(data, bytes): - raise CaptureStateError("compute log append data must be bytes") - lock: threading.RLock = self._capture_lock(invocation_id) - with lock: - writer: BinaryIO = self._active_writers(invocation_id)[stream] - writer.write(data) - writer.flush() - - def append_diagnostic(self, *, invocation_id: str, record: DiagnosticLog) -> None: - """Append one validated deterministic UTF-8 diagnostic JSON line.""" - - if not isinstance(record, DiagnosticLog): - raise CaptureStateError("diagnostic append requires DiagnosticLog") - encoded: bytes = (diagnostic_log_to_json(record) + "\n").encode("utf-8") - self.append( - invocation_id=invocation_id, - stream=ComputeLogStream.DIAGNOSTICS, - data=encoded, - ) - - def read( - self, - *, - invocation_id: str, - stream: ComputeLogStream, - cursor: ByteCursor = 0, - max_bytes: int = MAX_READ_BYTES, - ) -> ComputeLogReadChunk: - """Read at most the public byte limit from an exact stream offset.""" - - self._validate_capture_id(invocation_id) - if not isinstance(stream, ComputeLogStream): - raise CaptureStateError("stream must be a ComputeLogStream") - if type(cursor) is not int or cursor < 0: - raise InvalidComputeLogCursorError("cursor must be a nonnegative integer") - if type(max_bytes) is not int or max_bytes < 1 or max_bytes > MAX_READ_BYTES: - raise InvalidComputeLogLimitError( - f"max_bytes must be a positive integer no greater than {MAX_READ_BYTES}" - ) - capture_dir: Path = self._find_capture(invocation_id=invocation_id) - path: Path = capture_dir / _STREAM_FILES[stream] - marker: Path = capture_dir / COMPLETE_FILE_NAME - self._ensure_contained(path=path) - self._ensure_contained(path=marker) - complete: bool = marker.is_file() - if not path.exists(): - return ComputeLogReadChunk(data=b"", next_cursor=cursor, is_complete=complete) - size: int = path.stat().st_size - if cursor > size: - raise InvalidComputeLogCursorError( - f"cursor {cursor} is beyond current stream size {size}" - ) - with path.open("rb") as stream_file: - stream_file.seek(cursor) - data: bytes = stream_file.read(max_bytes) - next_cursor: int = stream_file.tell() - return ComputeLogReadChunk(data=data, next_cursor=next_cursor, is_complete=complete) - - def get_metadata(self, *, invocation_id: str) -> CaptureMetadata | FinalCaptureMetadata: - """Read the currently atomically published metadata document.""" - - capture_dir: Path = self._find_capture(invocation_id=invocation_id) - path: Path = capture_dir / METADATA_FILE_NAME - self._ensure_contained(path=path) - try: - metadata: CaptureMetadata | FinalCaptureMetadata = metadata_from_json( - path.read_text(encoding="utf-8") - ) - self._validate_metadata(metadata=metadata) - self._validate_directory_metadata(capture_dir=capture_dir, metadata=metadata) - return metadata - except OSError as error: - raise ComputeLogMetadataError(f"unable to read capture metadata: {error}") from error - - def get_byte_count(self, *, invocation_id: str, stream: ComputeLogStream) -> int: - """Return the current raw byte count for one stream.""" - - if not isinstance(stream, ComputeLogStream): - raise CaptureStateError("stream must be a ComputeLogStream") - counts: CaptureByteCounts = self.get_byte_counts(invocation_id=invocation_id) - if stream == ComputeLogStream.STDOUT: - return counts.stdout_bytes - if stream == ComputeLogStream.STDERR: - return counts.stderr_bytes - return counts.diagnostics_bytes - - def get_byte_counts(self, *, invocation_id: str) -> CaptureByteCounts: - """Return one lock-consistent snapshot of all stream sizes.""" - - self._validate_capture_id(invocation_id) - with self._state_lock: - lock: threading.RLock | None = self._locks.get(invocation_id) - if lock is None: - return self._read_byte_counts(invocation_id=invocation_id) - with lock: - return self._read_byte_counts(invocation_id=invocation_id) - - def is_complete(self, *, invocation_id: str) -> bool: - """Return capture-close evidence without asserting execution lifecycle status.""" - - capture_dir: Path = self._find_capture(invocation_id=invocation_id) - marker: Path = capture_dir / COMPLETE_FILE_NAME - self._ensure_contained(path=marker) - return marker.is_file() - - def mark_complete(self, metadata: FinalCaptureMetadata) -> None: - """Flush, fsync, and close streams before publishing final metadata and marker.""" - - self._ensure_open() - self._validate_metadata(metadata=metadata) - lock: threading.RLock = self._capture_lock(metadata.invocation_id) - with lock: - capture_dir: Path = self._find_capture(invocation_id=metadata.invocation_id) - initial_metadata: CaptureMetadata | FinalCaptureMetadata = self.get_metadata( - invocation_id=metadata.invocation_id - ) - if not isinstance(initial_metadata, CaptureMetadata) or ( - initial_metadata.started_at != metadata.started_at - or initial_metadata.capture_date != metadata.capture_date - or initial_metadata.command != metadata.command - or initial_metadata.project_dir != metadata.project_dir - or initial_metadata.target != metadata.target - or initial_metadata.run_id != metadata.run_id - ): - raise ComputeLogMetadataError( - "final metadata does not match capture start metadata" - ) - writers: dict[ComputeLogStream, BinaryIO] = self._active_writers(metadata.invocation_id) - actual_counts: dict[ComputeLogStream, int] = {} - for stream, writer in writers.items(): - writer.flush() - actual_counts[stream] = writer.tell() - if ( - metadata.stdout_bytes != actual_counts[ComputeLogStream.STDOUT] - or metadata.stderr_bytes != actual_counts[ComputeLogStream.STDERR] - or metadata.diagnostics_bytes != actual_counts[ComputeLogStream.DIAGNOSTICS] - ): - raise ComputeLogMetadataError("final byte counts do not match captured streams") - for writer in writers.values(): - os.fsync(writer.fileno()) - close_error: Exception | None = None - for writer in writers.values(): - try: - writer.close() - except Exception as error: - if close_error is None: - close_error = error - try: - if not writer.closed: - writer.close() - except Exception as retry_error: - if close_error is None: - close_error = retry_error - with self._state_lock: - self._writers.pop(metadata.invocation_id, None) - self._locks.pop(metadata.invocation_id, None) - if close_error is not None: - raise CaptureStateError("capture stream close failed") from close_error - self._write_metadata(capture_dir=capture_dir, metadata=metadata) - marker: Path = capture_dir / COMPLETE_FILE_NAME - with marker.open("xb") as marker_file: - marker_file.flush() - os.fsync(marker_file.fileno()) - self._fsync_directory(capture_dir) - - def delete(self, *, invocation_id: str) -> None: - """Delete a complete inactive capture without following symlinks.""" - - self._validate_capture_id(invocation_id) - with self._state_lock: - if invocation_id in self._writers: - raise CaptureStateError("active captures cannot be deleted") - capture_dir: Path = self._find_capture(invocation_id=invocation_id) - marker: Path = capture_dir / COMPLETE_FILE_NAME - self._ensure_contained(path=marker) - if not marker.is_file(): - raise CaptureStateError("incomplete captures cannot be deleted") - self._ensure_contained(path=capture_dir) - shutil.rmtree(capture_dir) - - def prune(self, retain_count: int | None = None) -> PruneResult: - """Retain the newest complete captures while never deleting incomplete captures.""" - - effective_count: int | None = ( - self._retention_count if retain_count is None else retain_count - ) - if effective_count is None: - inventory: CaptureInventory = self.inventory() - return PruneResult((), inventory.complete_count, inventory.incomplete_count) - if type(effective_count) is not int or effective_count < 0: - raise InvalidComputeLogLimitError("retain_count must be nonnegative or None") - lock_path: Path = self._root / PRUNE_LOCK_FILE_NAME - self._ensure_contained(path=lock_path) - with FileLock(lock_path): - inventory = self.inventory() - complete: list[CaptureInventoryItem] = [ - item for item in inventory.captures if item.is_complete - ] - complete.sort(key=self._completion_key, reverse=True) - deleted: list[str] = [] - for item in complete[effective_count:]: - capture_dir: Path = Path(item.path) - self._ensure_not_symlink(path=capture_dir) - self._ensure_contained(path=capture_dir) - marker: Path = capture_dir / COMPLETE_FILE_NAME - self._ensure_contained(path=marker) - if not marker.is_file(): - continue - self.delete(invocation_id=item.invocation_id) - deleted.append(item.invocation_id) - return PruneResult( - deleted_invocation_ids=tuple(deleted), - retained_complete_count=len(complete) - len(deleted), - retained_incomplete_count=inventory.incomplete_count, - ) - - def inventory(self) -> CaptureInventory: - """Inventory complete and abandoned captures without synthesizing lifecycle state.""" - - captures: list[CaptureInventoryItem] = [] - for date_dir in sorted(self._root.iterdir()): - if date_dir.name == PRUNE_LOCK_FILE_NAME: - continue - if not _DATE_PATTERN.fullmatch(date_dir.name): - continue - self._ensure_not_symlink(path=date_dir) - self._ensure_contained(path=date_dir) - if not date_dir.is_dir(): - continue - for capture_dir in sorted(date_dir.iterdir()): - if capture_dir.name.startswith("."): - continue - self._ensure_not_symlink(path=capture_dir) - self._ensure_contained(path=capture_dir) - if not capture_dir.is_dir(): - continue - metadata_path: Path = capture_dir / METADATA_FILE_NAME - marker_path: Path = capture_dir / COMPLETE_FILE_NAME - self._ensure_contained(path=metadata_path) - self._ensure_contained(path=marker_path) - metadata: CaptureMetadata | FinalCaptureMetadata = metadata_from_json( - metadata_path.read_text(encoding="utf-8") - ) - self._validate_metadata(metadata=metadata) - self._validate_directory_metadata(capture_dir=capture_dir, metadata=metadata) - captures.append( - CaptureInventoryItem( - invocation_id=metadata.invocation_id, - capture_date=metadata.capture_date, - path=str(capture_dir), - is_complete=marker_path.is_file(), - metadata=metadata, - ) - ) - captures.sort(key=lambda item: (item.capture_date, item.invocation_id)) - complete_count: int = sum(item.is_complete for item in captures) - return CaptureInventory( - captures=tuple(captures), - complete_count=complete_count, - incomplete_count=len(captures) - complete_count, - ) - - def close(self) -> None: - """Close active writers, deliberately leaving their captures incomplete.""" - - with self._state_lock: - captures: tuple[tuple[threading.RLock, dict[ComputeLogStream, BinaryIO]], ...] = tuple( - (self._locks[invocation_id], writers) - for invocation_id, writers in self._writers.items() - ) - self._writers.clear() - self._locks.clear() - self._closed = True - first_error: Exception | None = None - for lock, writers in captures: - with lock: - for writer in writers.values(): - try: - if not writer.closed: - writer.flush() - except Exception as error: - if first_error is None: - first_error = error - try: - if not writer.closed: - writer.close() - except Exception as error: - if first_error is None: - first_error = error - try: - if not writer.closed: - writer.close() - except Exception as retry_error: - if first_error is None: - first_error = retry_error - if first_error is not None: - raise CaptureStateError("compute log writer cleanup failed") from first_error - - def dispose(self) -> None: - """Dispose process-local resources without deleting retained logs.""" - - self.close() - - def __enter__(self) -> Self: - return self - - def __exit__( - self, - exception_type: type[BaseException] | None, - exception: BaseException | None, - traceback: TracebackType | None, - ) -> bool | None: - self.close() - return None - - @overload - def _find_capture(self, *, invocation_id: str, missing_ok: Literal[False] = False) -> Path: ... - - @overload - def _find_capture(self, *, invocation_id: str, missing_ok: Literal[True]) -> Path | None: ... - - def _find_capture(self, *, invocation_id: str, missing_ok: bool = False) -> Path | None: - self._validate_capture_id(invocation_id) - matches: list[Path] = [] - for date_dir in self._root.iterdir(): - if not _DATE_PATTERN.fullmatch(date_dir.name): - continue - self._ensure_not_symlink(path=date_dir) - self._ensure_contained(path=date_dir) - if not date_dir.is_dir(): - continue - candidate: Path = date_dir / invocation_id - if candidate.is_symlink(): - raise ComputeLogPathError(f"compute log path cannot be a symlink: {candidate}") - if candidate.exists(): - self._ensure_contained(path=candidate) - matches.append(candidate) - if len(matches) > 1: - raise ComputeLogMetadataError(f"duplicate capture identity {invocation_id!r}") - if matches: - return matches[0] - if missing_ok: - return None - raise CaptureNotFoundError(f"capture not found for invocation {invocation_id!r}") - - def _write_metadata( - self, *, capture_dir: Path, metadata: CaptureMetadata | FinalCaptureMetadata - ) -> None: - temporary: Path = capture_dir / f".{METADATA_FILE_NAME}.{uuid4().hex}.tmp" - destination: Path = capture_dir / METADATA_FILE_NAME - self._ensure_contained(path=temporary) - try: - with temporary.open("x", encoding="utf-8", newline="") as metadata_file: - metadata_file.write(metadata_to_json(metadata)) - metadata_file.flush() - os.fsync(metadata_file.fileno()) - os.replace(temporary, destination) - self._fsync_directory(capture_dir) - finally: - temporary.unlink(missing_ok=True) - - def _capture_lock(self, invocation_id: str) -> threading.RLock: - with self._state_lock: - try: - return self._locks[invocation_id] - except KeyError as error: - raise CaptureStateError("capture is not active in this storage instance") from error - - def _active_writers(self, invocation_id: str) -> dict[ComputeLogStream, BinaryIO]: - try: - return self._writers[invocation_id] - except KeyError as error: - raise CaptureStateError("capture is not active in this storage instance") from error - - def _read_byte_counts(self, *, invocation_id: str) -> CaptureByteCounts: - capture_dir: Path = self._find_capture(invocation_id=invocation_id) - sizes: dict[ComputeLogStream, int] = {} - for stream, file_name in _STREAM_FILES.items(): - path: Path = capture_dir / file_name - self._ensure_contained(path=path) - sizes[stream] = path.stat().st_size if path.exists() else 0 - return CaptureByteCounts( - stdout_bytes=sizes[ComputeLogStream.STDOUT], - stderr_bytes=sizes[ComputeLogStream.STDERR], - diagnostics_bytes=sizes[ComputeLogStream.DIAGNOSTICS], - ) - - def _validate_directory_metadata( - self, - *, - capture_dir: Path, - metadata: CaptureMetadata | FinalCaptureMetadata, - ) -> None: - if metadata.invocation_id != capture_dir.name: - raise ComputeLogMetadataError("metadata invocation_id does not match capture directory") - if metadata.capture_date != capture_dir.parent.name: - raise ComputeLogMetadataError("metadata capture_date does not match date directory") - - def _validate_metadata(self, metadata: CaptureMetadata | FinalCaptureMetadata) -> None: - self._validate_capture_id(metadata.invocation_id) - if metadata.format_version != COMPUTE_LOG_FORMAT_VERSION: - raise ComputeLogMetadataError("unsupported compute log format version") - if not metadata.command.strip() or not metadata.project_dir.strip(): - raise ComputeLogMetadataError("command and project_dir must be non-empty") - if metadata.started_at.tzinfo is None or metadata.started_at.utcoffset() is None: - raise ComputeLogMetadataError("started_at must be timezone-aware") - expected_date: str = metadata.started_at.astimezone(UTC).date().isoformat() - if metadata.capture_date != expected_date: - raise ComputeLogMetadataError("capture_date must derive from UTC started_at") - if isinstance(metadata, FinalCaptureMetadata): - if metadata.completed_at.tzinfo is None or metadata.completed_at < metadata.started_at: - raise ComputeLogMetadataError( - "completed_at must be aware and not precede started_at" - ) - if type(metadata.exit_code) is not int or any( - type(count) is not int - for count in ( - metadata.stdout_bytes, - metadata.stderr_bytes, - metadata.diagnostics_bytes, - ) - ): - raise ComputeLogMetadataError("exit_code and byte counts must be integers") - if ( - min( - metadata.stdout_bytes, - metadata.stderr_bytes, - metadata.diagnostics_bytes, - ) - < 0 - ): - raise ComputeLogMetadataError("final byte counts must be nonnegative") - - def _validate_capture_id(self, invocation_id: str) -> None: - if not isinstance(invocation_id, str) or not _CAPTURE_ID_PATTERN.fullmatch(invocation_id): - raise InvalidCaptureIdError( - "invocation_id must contain only ASCII letters, digits, underscores, or hyphens" - ) - - def _ensure_contained(self, *, path: Path) -> None: - resolved: Path = path.resolve(strict=False) - if not resolved.is_relative_to(self._root): - raise ComputeLogPathError(f"compute log path escapes configured root: {path}") - - @staticmethod - def _ensure_not_symlink(*, path: Path) -> None: - if path.is_symlink(): - raise ComputeLogPathError(f"compute log path cannot be a symlink: {path}") - - def _ensure_open(self) -> None: - if self._closed: - raise CaptureStateError("compute log storage is closed") - - @staticmethod - def _completion_key(item: CaptureInventoryItem) -> tuple[object, str]: - metadata: CaptureMetadata | FinalCaptureMetadata = item.metadata - completed_at: object = ( - metadata.completed_at - if isinstance(metadata, FinalCaptureMetadata) - else metadata.started_at - ) - return completed_at, item.invocation_id - - @staticmethod - def _fsync_directory(path: Path) -> None: - descriptor: int = os.open(path, os.O_RDONLY) - try: - os.fsync(descriptor) - finally: - os.close(descriptor) diff --git a/src/sqlbuild/runtime/compute_logs/classes/scoped_compute_log_capture.py b/src/sqlbuild/runtime/compute_logs/classes/scoped_compute_log_capture.py index 5025e62098..9728eae35a 100644 --- a/src/sqlbuild/runtime/compute_logs/classes/scoped_compute_log_capture.py +++ b/src/sqlbuild/runtime/compute_logs/classes/scoped_compute_log_capture.py @@ -7,11 +7,9 @@ from datetime import UTC, datetime from typing import Any +from sqlbuild._native import LocalFilesystemComputeLogStorage from sqlbuild.diagnostics.classes.invocation_diagnostic_routing import InvocationDiagnosticRouting from sqlbuild.diagnostics.models import DiagnosticRoutingOptions -from sqlbuild.runtime.compute_logs.classes.local_filesystem_compute_log_storage import ( - LocalFilesystemComputeLogStorage, -) from sqlbuild.runtime.compute_logs.classes.text_tee import TextComputeLogTee from sqlbuild.runtime.compute_logs.models import ( CaptureByteCounts, diff --git a/src/sqlbuild/runtime/compute_logs/main/__init__.py b/src/sqlbuild/runtime/compute_logs/main/__init__.py index 7bd0dd6e3e..2a1b02b821 100644 --- a/src/sqlbuild/runtime/compute_logs/main/__init__.py +++ b/src/sqlbuild/runtime/compute_logs/main/__init__.py @@ -1 +1 @@ -"""Compute log construction entrypoints.""" +"""Compute log capture entry points.""" diff --git a/src/sqlbuild/runtime/compute_logs/main/open_local_compute_log_storage.py b/src/sqlbuild/runtime/compute_logs/main/open_local_compute_log_storage.py deleted file mode 100644 index 4d8744e67d..0000000000 --- a/src/sqlbuild/runtime/compute_logs/main/open_local_compute_log_storage.py +++ /dev/null @@ -1,23 +0,0 @@ -"""Local compute log storage construction entrypoint.""" - -from pathlib import Path - -from sqlbuild.runtime.compute_logs.classes.local_filesystem_compute_log_storage import ( - LocalFilesystemComputeLogStorage, -) -from sqlbuild.runtime.compute_logs.constants import DEFAULT_RETENTION_COUNT - - -def open_local_compute_log_storage( - *, - project_dir: Path, - root: Path | None = None, - retention_count: int | None = DEFAULT_RETENTION_COUNT, -) -> LocalFilesystemComputeLogStorage: - """Open the supported project-local compute log backend.""" - - return LocalFilesystemComputeLogStorage( - project_dir=project_dir, - root=root, - retention_count=retention_count, - ) diff --git a/src/sqlbuild/runtime/compute_logs/main/run_with_compute_log_capture.py b/src/sqlbuild/runtime/compute_logs/main/run_with_compute_log_capture.py new file mode 100644 index 0000000000..503ce48c65 --- /dev/null +++ b/src/sqlbuild/runtime/compute_logs/main/run_with_compute_log_capture.py @@ -0,0 +1,29 @@ +"""Run one operation with its process streams and diagnostics captured.""" + +from collections.abc import Callable + +from sqlbuild._native import LocalFilesystemComputeLogStorage +from sqlbuild.diagnostics.models import DiagnosticRoutingOptions +from sqlbuild.runtime.compute_logs.classes.scoped_compute_log_capture import ( + ScopedComputeLogCapture, +) +from sqlbuild.runtime.compute_logs.models import CaptureMetadata + + +def run_with_compute_log_capture( + *, + storage: LocalFilesystemComputeLogStorage, + metadata: CaptureMetadata, + failure_callback: Callable[[Exception, str], None], + routing_options: DiagnosticRoutingOptions, + operation: Callable[[], int], +) -> int: + """Capture streams and diagnostics around `operation`, preserving its outcome.""" + + capture: ScopedComputeLogCapture = ScopedComputeLogCapture( + storage=storage, + metadata=metadata, + failure_callback=failure_callback, + routing_options=routing_options, + ) + return capture.run(operation=operation) diff --git a/tests/unit/src/sqlbuild/runtime/compute_logs/classes/helpers.py b/tests/unit/src/sqlbuild/runtime/compute_logs/classes/helpers.py index 7601b40631..23d1a46d00 100644 --- a/tests/unit/src/sqlbuild/runtime/compute_logs/classes/helpers.py +++ b/tests/unit/src/sqlbuild/runtime/compute_logs/classes/helpers.py @@ -1,8 +1,6 @@ import logging -from collections.abc import Callable, Iterator from datetime import UTC, datetime from pathlib import Path -from typing import Any, BinaryIO from sqlbuild.compute_logs import ( COMPUTE_LOG_FORMAT_VERSION, @@ -28,6 +26,10 @@ def flush(self) -> None: return None +class PatchableComputeLogStorage(LocalFilesystemComputeLogStorage): + """Native storage whose instance methods tests may replace.""" + + class InvalidLogMessage: def __str__(self) -> str: raise ValueError("controlled invalid message") @@ -47,30 +49,6 @@ def close(self) -> None: super().close() -class FailOnceBinaryWriter: - def __init__(self, writer: BinaryIO) -> None: - self._writer: BinaryIO = writer - self.close_calls: int = 0 - self._close_actions: Iterator[Callable[[], None]] = iter( - (self._fail_first_close, writer.close) - ) - - def close(self) -> None: - self.close_calls += 1 - next(self._close_actions)() - - @property - def closed(self) -> bool: - return self._writer.closed - - def __getattr__(self, name: str) -> Any: - return getattr(self._writer, name) - - @staticmethod - def _fail_first_close() -> None: - raise OSError("controlled first close failure") - - def build_capture_metadata( *, project_dir: Path, invocation_id: str, started_at: datetime ) -> CaptureMetadata: diff --git a/tests/unit/src/sqlbuild/runtime/compute_logs/classes/test_scoped_compute_log_capture.py b/tests/unit/src/sqlbuild/runtime/compute_logs/classes/test_scoped_compute_log_capture.py index 086d9956e6..f35e26dba5 100644 --- a/tests/unit/src/sqlbuild/runtime/compute_logs/classes/test_scoped_compute_log_capture.py +++ b/tests/unit/src/sqlbuild/runtime/compute_logs/classes/test_scoped_compute_log_capture.py @@ -10,6 +10,7 @@ import pytest +from sqlbuild._native import LocalFilesystemComputeLogStorage from sqlbuild.compute_logs import ( CaptureMetadata, CaptureStateError, @@ -24,9 +25,6 @@ from sqlbuild.runtime.compute_logs.classes.diagnostic_log_handler import ( ComputeDiagnosticLogHandler, ) -from sqlbuild.runtime.compute_logs.classes.local_filesystem_compute_log_storage import ( - LocalFilesystemComputeLogStorage, -) from sqlbuild.runtime.compute_logs.classes.scoped_compute_log_capture import ( ScopedComputeLogCapture, ) @@ -40,6 +38,7 @@ from tests.unit.src.sqlbuild.runtime.compute_logs.classes.helpers import ( AccessFailureTextSink, HostRecordingLogHandler, + PatchableComputeLogStorage, build_capture_metadata, ) @@ -226,9 +225,7 @@ def test_given_handler_removal_failure_when_cleaning_then_fallback_prevents_hand def test_given_finalization_failure_when_operation_returns_then_result_is_not_replaced( tmp_path: Path, test_case: CaptureOutcomeTestCase ) -> None: - storage: LocalFilesystemComputeLogStorage = LocalFilesystemComputeLogStorage( - project_dir=tmp_path - ) + storage: LocalFilesystemComputeLogStorage = PatchableComputeLogStorage(project_dir=tmp_path) metadata: CaptureMetadata = build_capture_metadata( project_dir=tmp_path, invocation_id="finalize_failure", started_at=datetime.now(UTC) ) @@ -261,9 +258,7 @@ def test_given_finalization_failure_when_operation_returns_then_result_is_not_re def test_given_storage_close_failure_when_cleanup_runs_then_result_and_writer_cleanup_are_preserved( tmp_path: Path, test_case: CaptureOutcomeTestCase ) -> None: - storage: LocalFilesystemComputeLogStorage = LocalFilesystemComputeLogStorage( - project_dir=tmp_path - ) + storage: LocalFilesystemComputeLogStorage = PatchableComputeLogStorage(project_dir=tmp_path) metadata: CaptureMetadata = build_capture_metadata( project_dir=tmp_path, invocation_id="close_failure", started_at=datetime.now(UTC) ) @@ -556,9 +551,7 @@ def emit_records() -> int: def test_given_storage_append_failure_when_operation_writes_then_failure_reaches_callback( tmp_path: Path, test_case: CaptureWriteFailureTestCase ) -> None: - storage: LocalFilesystemComputeLogStorage = LocalFilesystemComputeLogStorage( - project_dir=tmp_path - ) + storage: LocalFilesystemComputeLogStorage = PatchableComputeLogStorage(project_dir=tmp_path) metadata: CaptureMetadata = build_capture_metadata( project_dir=tmp_path, invocation_id="append_failure", started_at=datetime.now(UTC) ) diff --git a/tests/unit/src/sqlbuild/runtime/compute_logs/classes/test_storage_recovery_and_paths.py b/tests/unit/src/sqlbuild/runtime/compute_logs/classes/test_storage_recovery_and_paths.py index 0ccd344b94..a983855769 100644 --- a/tests/unit/src/sqlbuild/runtime/compute_logs/classes/test_storage_recovery_and_paths.py +++ b/tests/unit/src/sqlbuild/runtime/compute_logs/classes/test_storage_recovery_and_paths.py @@ -3,15 +3,12 @@ from dataclasses import replace from datetime import UTC, datetime, timedelta from pathlib import Path -from typing import BinaryIO, cast -from unittest.mock import Mock, patch import pytest from sqlbuild.compute_logs import ( CaptureInventory, CaptureMetadata, - CaptureStateError, ComputeLogMetadataError, ComputeLogPathError, ComputeLogReadChunk, @@ -26,7 +23,6 @@ RetentionTestCase, ) from tests.unit.src.sqlbuild.runtime.compute_logs.classes.helpers import ( - FailOnceBinaryWriter, build_capture_metadata, build_final_metadata, ) @@ -86,99 +82,6 @@ def test_given_mismatched_completion_counts_when_retrying_then_capture_remains_r storage.close() -@pytest.mark.parametrize( - "test_case", - ( - RecoveryTestCase( - description="writer close failure deactivates capture and closes retryable handle", - initial_bytes=b"captured", - retry_bytes=b"", - expected_bytes=b"captured", - expected_complete=False, - ), - ), - ids=lambda case: case.description, -) -def test_given_writer_close_failure_when_completing_then_incomplete_state_remains_coherent( - tmp_path: Path, test_case: RecoveryTestCase -) -> None: - storage: LocalFilesystemComputeLogStorage = LocalFilesystemComputeLogStorage( - project_dir=tmp_path - ) - initial: CaptureMetadata = build_capture_metadata( - project_dir=tmp_path, invocation_id="writer_close_failure", started_at=datetime.now(UTC) - ) - storage.start_capture(initial) - storage.append( - invocation_id=initial.invocation_id, - stream=ComputeLogStream.STDOUT, - data=test_case.initial_bytes, - ) - original_writer: BinaryIO = storage._writers[initial.invocation_id][ComputeLogStream.STDOUT] - failing_writer: FailOnceBinaryWriter = FailOnceBinaryWriter(original_writer) - storage._writers[initial.invocation_id][ComputeLogStream.STDOUT] = cast( - BinaryIO, failing_writer - ) - final_metadata: FinalCaptureMetadata = build_final_metadata( - storage=storage, initial=initial, completed_at=datetime.now(UTC) - ) - - with pytest.raises(CaptureStateError, match="stream close failed"): - storage.mark_complete(final_metadata) - retained: ComputeLogReadChunk = storage.read( - invocation_id=initial.invocation_id, stream=ComputeLogStream.STDOUT - ) - - assert retained.data == test_case.expected_bytes - assert retained.is_complete is test_case.expected_complete - assert failing_writer.closed is True - assert failing_writer.close_calls == 2 - storage.close() - - -@pytest.mark.parametrize( - "test_case", - ( - RecoveryTestCase( - description="failed staged publication leaves no visible capture", - initial_bytes=b"", - retry_bytes=b"", - expected_bytes=b"", - expected_complete=False, - ), - ), - ids=lambda case: case.description, -) -def test_given_initial_metadata_failure_when_publishing_then_staging_directory_is_cleaned( - tmp_path: Path, test_case: RecoveryTestCase -) -> None: - storage: LocalFilesystemComputeLogStorage = LocalFilesystemComputeLogStorage( - project_dir=tmp_path - ) - initial: CaptureMetadata = build_capture_metadata( - project_dir=tmp_path, invocation_id="staged_failure", started_at=datetime.now(UTC) - ) - - with ( - patch.object( - LocalFilesystemComputeLogStorage, - "_write_metadata", - Mock(side_effect=OSError("controlled metadata failure")), - ), - pytest.raises(OSError, match="controlled metadata failure"), - ): - storage.start_capture(initial) - date_dir: Path = storage.root / initial.capture_date - entries: tuple[Path, ...] = tuple(date_dir.iterdir()) - inventory: CaptureInventory = storage.inventory() - - assert entries == () - assert inventory.captures == () - assert test_case.expected_bytes == b"" - assert test_case.expected_complete is False - storage.close() - - @pytest.mark.parametrize( "test_case", ( From bf09e94ecdb23e150df19e2cc6673db232394411 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:55:50 +0100 Subject: [PATCH 09/12] fix: build compute log storage errors after releasing the state lock Constructing a Python exception imports the exceptions module, which can release the GIL. Doing that while holding the storage state mutex let another thread take the GIL and block on the mutex, deadlocking a late diagnostic append against storage close. Every Python call now happens after the guard is dropped. --- .../bindings/_helpers/compute_logs/storage.rs | 81 +++-- .../output_capture/_helpers/__init__.py | 1 - .../runtime/output_capture/_helpers/json.py | 35 -- .../runtime/output_capture/_helpers/scope.py | 33 -- .../runtime/output_capture/_helpers/text.py | 88 ----- .../output_capture/classes/dispatcher.py | 304 ------------------ .../main/command_output_context.py | 14 - .../main/command_output_from_json.py | 83 ----- .../main/command_output_to_json.py | 47 --- .../main/configured_output_capture_scope.py} | 0 .../main/current_output_capture_context.py | 12 - .../output_capture/main/invocation_context.py | 61 ---- 12 files changed, 55 insertions(+), 704 deletions(-) delete mode 100644 src/sqlbuild/runtime/output_capture/_helpers/__init__.py delete mode 100644 src/sqlbuild/runtime/output_capture/_helpers/json.py delete mode 100644 src/sqlbuild/runtime/output_capture/_helpers/scope.py delete mode 100644 src/sqlbuild/runtime/output_capture/_helpers/text.py delete mode 100644 src/sqlbuild/runtime/output_capture/classes/dispatcher.py delete mode 100644 src/sqlbuild/runtime/output_capture/main/command_output_context.py delete mode 100644 src/sqlbuild/runtime/output_capture/main/command_output_from_json.py delete mode 100644 src/sqlbuild/runtime/output_capture/main/command_output_to_json.py rename src/sqlbuild/{cli/commands/main/entrypoint/_dispatch_with_output_capture.py => runtime/output_capture/main/configured_output_capture_scope.py} (100%) delete mode 100644 src/sqlbuild/runtime/output_capture/main/current_output_capture_context.py delete mode 100644 src/sqlbuild/runtime/output_capture/main/invocation_context.py diff --git a/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/storage.rs b/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/storage.rs index c33ecdee32..a722957f46 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/storage.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/storage.rs @@ -26,6 +26,7 @@ use sqlbuild_cli::compute_logs::main::read_capture_document::read_capture_docume use sqlbuild_cli::compute_logs::main::read_stream_chunk::read_stream_chunk; use sqlbuild_cli::compute_logs::main::stream_byte_counts::stream_byte_counts; use sqlbuild_cli::compute_logs::main::validate_capture_id::validate_capture_id; +use sqlbuild_cli::compute_logs::models::CaptureDocument; use crate::bindings::_helpers::compile_driver::python_calls::{ KeywordArguments, call_with, imported, @@ -55,9 +56,9 @@ pub(crate) struct LocalFilesystemComputeLogStorage { } /// The stdout, stderr and diagnostics byte counts of open streams, as `tell()` reports them. -fn stream_positions(streams: &mut [File; 3]) -> PyResult<[i64; 3]> { +fn stream_positions(streams: &[File; 3]) -> std::io::Result<[i64; 3]> { let mut counts: [i64; 3] = [0; 3]; - for (count, stream) in counts.iter_mut().zip(streams.iter_mut()) { + for (count, mut stream) in counts.iter_mut().zip(streams.iter()) { *count = i64::try_from(stream.stream_position()?).unwrap_or(i64::MAX); } Ok(counts) @@ -108,12 +109,35 @@ impl LocalFilesystemComputeLogStorage { stream: usize, data: &[u8], ) -> PyResult<()> { - let mut state = self.locked()?; - let Some(writers) = state.writers.get_mut(invocation_id) else { - return Err(not_active(py)); + let written: Option> = { + let mut state = self.locked()?; + state + .writers + .get_mut(invocation_id) + .map(|writers| writers[stream].write_all(data)) }; - writers[stream].write_all(data)?; - Ok(()) + match written { + Some(result) => Ok(result?), + None => Err(not_active(py)), + } + } + + /// The capture directory once final metadata matches the open, synced streams. + fn checked_completion( + &self, + state: &StorageState, + document: &CaptureDocument, + ) -> Result { + let streams = state.writers.get(&document.invocation_id); + let counts = match streams { + Some(streams) => Some(stream_positions(streams)?), + None => None, + }; + let directory = check_completion(&self.root, document, counts)?; + for stream in streams.into_iter().flatten() { + stream.sync_all()?; + } + Ok(directory) } fn byte_counts(&self, py: Python<'_>, invocation_id: &Bound<'_, PyAny>) -> PyResult<[u64; 3]> { @@ -153,10 +177,17 @@ impl LocalFilesystemComputeLogStorage { fn start_capture(&self, py: Python<'_>, metadata: &Bound<'_, PyAny>) -> PyResult<()> { self.ensure_open(py)?; let document = document(metadata)?; - let mut state = self.locked()?; - let writers = checked(py, publish_capture(&self.root, &document))?; - state.writers.insert(document.invocation_id, writers); - Ok(()) + let published = { + let mut state = self.locked()?; + publish_capture(&self.root, &document).map(|writers| { + drop( + state + .writers + .insert(document.invocation_id.clone(), writers), + ) + }) + }; + checked(py, published) } /// Append exact raw bytes and make them visible to concurrent readers. @@ -311,22 +342,20 @@ impl LocalFilesystemComputeLogStorage { fn mark_complete(&self, py: Python<'_>, metadata: &Bound<'_, PyAny>) -> PyResult<()> { self.ensure_open(py)?; let document = document(metadata)?; - let mut state = self.locked()?; - let counts = match state.writers.get_mut(&document.invocation_id) { - Some(streams) => Some(stream_positions(streams)?), - None => None, - }; - let directory = checked(py, check_completion(&self.root, &document, counts))?; - let Some(streams) = state.writers.get(&document.invocation_id) else { - return Err(not_active(py)); - }; - for stream in streams { - stream.sync_all()?; - } - let Some(streams) = state.writers.remove(&document.invocation_id) else { - return Err(not_active(py)); + let checked_streams = { + let mut state = self.locked()?; + self.checked_completion(&state, &document) + .and_then( + |directory| match state.writers.remove(&document.invocation_id) { + Some(streams) => Ok((directory, streams)), + None => Err(ComputeLogError::rejected( + ComputeLogErrorKind::State, + "capture is not active in this storage instance", + )), + }, + ) }; - drop(state); + let (directory, streams) = checked(py, checked_streams)?; checked(py, complete_capture(&directory, &document, streams)) } diff --git a/src/sqlbuild/runtime/output_capture/_helpers/__init__.py b/src/sqlbuild/runtime/output_capture/_helpers/__init__.py deleted file mode 100644 index a2b7d82516..0000000000 --- a/src/sqlbuild/runtime/output_capture/_helpers/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Output capture implementation helpers.""" diff --git a/src/sqlbuild/runtime/output_capture/_helpers/json.py b/src/sqlbuild/runtime/output_capture/_helpers/json.py deleted file mode 100644 index c66a12daa6..0000000000 --- a/src/sqlbuild/runtime/output_capture/_helpers/json.py +++ /dev/null @@ -1,35 +0,0 @@ -"""Command-output JSON compatibility helpers.""" - -from __future__ import annotations - -import math -from collections.abc import Mapping -from types import MappingProxyType - -from sqlbuild.runtime.output_capture.exceptions import CommandOutputValidationError - - -def freeze_command_output_json(*, value: object, path: str) -> object: - """Validate and recursively freeze a JSON-compatible value.""" - - if value is None or isinstance(value, (str, bool, int)): - return value - if isinstance(value, float): - if not math.isfinite(value): - raise CommandOutputValidationError(f"{path} must not contain NaN or infinity") - return value - if isinstance(value, Mapping): - frozen: dict[str, object] = {} - for key, item in value.items(): - if not isinstance(key, str): - raise CommandOutputValidationError(f"{path} keys must be strings") - frozen[key] = freeze_command_output_json(value=item, path=f"{path}.{key}") - return MappingProxyType(frozen) - if isinstance(value, (list, tuple)): - return tuple( - freeze_command_output_json(value=item, path=f"{path}[{index}]") - for index, item in enumerate(value) - ) - raise CommandOutputValidationError( - f"{path} contains non-JSON value of type {type(value).__name__}" - ) diff --git a/src/sqlbuild/runtime/output_capture/_helpers/scope.py b/src/sqlbuild/runtime/output_capture/_helpers/scope.py deleted file mode 100644 index ec02bf4864..0000000000 --- a/src/sqlbuild/runtime/output_capture/_helpers/scope.py +++ /dev/null @@ -1,33 +0,0 @@ -"""Integration-installed output exporter command context.""" - -from __future__ import annotations - -from collections.abc import Iterator, Mapping -from contextlib import contextmanager -from contextvars import ContextVar, Token - -from sqlbuild.runtime.output_capture.models import OutputCaptureContext - -_CURRENT_CONTEXT: ContextVar[OutputCaptureContext | None] = ContextVar( - "sqlbuild_output_capture_context", default=None -) - - -@contextmanager -def output_capture_context(*, external_context: Mapping[str, object]) -> Iterator[None]: - """Attach integration context without interpreting its contents.""" - - context: OutputCaptureContext = OutputCaptureContext( - external_context=dict(external_context), - ) - token: Token[OutputCaptureContext | None] = _CURRENT_CONTEXT.set(context) - try: - yield - finally: - _CURRENT_CONTEXT.reset(token) - - -def current_output_capture_context() -> OutputCaptureContext | None: - """Return optional integration context for the active command.""" - - return _CURRENT_CONTEXT.get() diff --git a/src/sqlbuild/runtime/output_capture/_helpers/text.py b/src/sqlbuild/runtime/output_capture/_helpers/text.py deleted file mode 100644 index 69e096b9b6..0000000000 --- a/src/sqlbuild/runtime/output_capture/_helpers/text.py +++ /dev/null @@ -1,88 +0,0 @@ -"""Plain-text normalization and deterministic chunking.""" - -from __future__ import annotations - -import re - -from sqlbuild.runtime.output_capture.constants import ( - UTF8_CONTINUATION_BYTE_MASK, - UTF8_CONTINUATION_BYTE_TAG, -) - -_ANSI_ESCAPE: re.Pattern[str] = re.compile( - r"(?:\x1b\][^\x07\x1b]*(?:\x07|\x1b\\)|\x1b\[[0-?]*[ -/]*[@-~])" -) - - -def strip_ansi(text: str) -> str: - """Remove terminal control sequences from only the exported copy.""" - - return _ANSI_ESCAPE.sub("", text) - - -def chunk_text(*, text: str, max_bytes: int) -> tuple[str, ...]: - """Split text deterministically without splitting Unicode code points.""" - - if text.isascii() and max_bytes > 0: - if len(text) <= max_bytes: - return (text,) - return tuple(text[index : index + max_bytes] for index in range(0, len(text), max_bytes)) - try: - encoded: bytes = text.encode("utf-8") - except UnicodeEncodeError: - return _chunk_text_by_character(text=text, max_bytes=max_bytes) - if len(encoded) <= max_bytes: - return (text,) - return _chunk_utf8(encoded=encoded, max_bytes=max_bytes) - - -def _chunk_utf8(*, encoded: bytes, max_bytes: int) -> tuple[str, ...]: - """Split valid UTF-8 greedily at the last code point boundary within each limit.""" - - chunks: list[str] = [] - total: int = len(encoded) - start: int = 0 - while start < total: - end: int = start + max_bytes - if end >= total: - end = total - else: - while ( - end > start - and encoded[end] & UTF8_CONTINUATION_BYTE_MASK == UTF8_CONTINUATION_BYTE_TAG - ): - end -= 1 - if end <= start: - end = start + 1 - while ( - end < total - and encoded[end] & UTF8_CONTINUATION_BYTE_MASK == UTF8_CONTINUATION_BYTE_TAG - ): - end += 1 - chunks.append(encoded[start:end].decode("utf-8")) - start = end - if not chunks: - chunks.append("") - return tuple(chunks) - - -def _chunk_text_by_character(*, text: str, max_bytes: int) -> tuple[str, ...]: - """Split text containing lone surrogates, which surrogateescape maps to single bytes.""" - - if len(text) <= max_bytes and len(text.encode("utf-8", "surrogateescape")) <= max_bytes: - return (text,) - - chunks: list[str] = [] - current: list[str] = [] - current_bytes: int = 0 - for character in text: - character_bytes: int = len(character.encode("utf-8", "surrogateescape")) - if current and current_bytes + character_bytes > max_bytes: - chunks.append("".join(current)) - current = [] - current_bytes = 0 - current.append(character) - current_bytes += character_bytes - if current or not chunks: - chunks.append("".join(current)) - return tuple(chunks) diff --git a/src/sqlbuild/runtime/output_capture/classes/dispatcher.py b/src/sqlbuild/runtime/output_capture/classes/dispatcher.py deleted file mode 100644 index 1c6f5a891f..0000000000 --- a/src/sqlbuild/runtime/output_capture/classes/dispatcher.py +++ /dev/null @@ -1,304 +0,0 @@ -"""Bounded asynchronous delivery of captured output batches.""" - -from __future__ import annotations - -import threading -import time -from collections import deque -from collections.abc import Callable, Mapping -from datetime import UTC, datetime -from typing import cast - -from sqlbuild.runtime.output_capture._helpers.json import freeze_command_output_json -from sqlbuild.runtime.output_capture._helpers.text import chunk_text, strip_ansi -from sqlbuild.runtime.output_capture.constants import ( - COMMAND_OUTPUT_LOSS_RECORD_TYPE, - DEFAULT_OUTPUT_BATCH_SIZE, - DEFAULT_OUTPUT_FLUSH_INTERVAL_SECONDS, - DEFAULT_OUTPUT_MAX_RECORD_BYTES, - DEFAULT_OUTPUT_QUEUE_CAPACITY, - DEFAULT_OUTPUT_SHUTDOWN_TIMEOUT_SECONDS, - MIN_OUTPUT_RECORD_BYTES, -) -from sqlbuild.runtime.output_capture.exceptions import OutputCaptureInputError -from sqlbuild.runtime.output_capture.models import ( - CommandOutputCaptureSummary, - CommandOutputRecord, -) -from sqlbuild.runtime.output_capture.types import ( - CommandOutputBatchExporter, - CommandOutputStream, - OutputRecordPriority, -) - - -class OutputCaptureDispatcher: - """Turn stream fragments into bounded text chunks and export them off execution threads.""" - - def __init__( - self, - *, - exporter: object, - invocation_id: str, - run_id: str | None = None, - external_context: Mapping[str, object] | None = None, - queue_capacity: int = DEFAULT_OUTPUT_QUEUE_CAPACITY, - batch_size: int = DEFAULT_OUTPUT_BATCH_SIZE, - max_record_bytes: int = DEFAULT_OUTPUT_MAX_RECORD_BYTES, - shutdown_timeout_seconds: float = DEFAULT_OUTPUT_SHUTDOWN_TIMEOUT_SECONDS, - clock: Callable[[], datetime] | None = None, - failure_callback: Callable[[BaseException], object] | None = None, - ) -> None: - if queue_capacity < 1 or batch_size < 1: - raise OutputCaptureInputError( - "output capture queue capacity and batch size must be positive" - ) - if max_record_bytes < MIN_OUTPUT_RECORD_BYTES or shutdown_timeout_seconds < 0: - raise OutputCaptureInputError("output capture limits are invalid") - self._exporter: object = exporter - self._invocation_id: str = invocation_id - self._run_id: str | None = run_id - frozen_context: object = freeze_command_output_json( - value=dict(external_context or {}), - path="external_context", - ) - if not isinstance(frozen_context, Mapping): - raise OutputCaptureInputError("output capture external_context must be a mapping") - self._context: Mapping[str, object] = cast(Mapping[str, object], frozen_context) - self._capacity: int = queue_capacity - self._batch_size: int = batch_size - self._max_record_bytes: int = max_record_bytes - self._flush_interval_seconds: float = DEFAULT_OUTPUT_FLUSH_INTERVAL_SECONDS - self._shutdown_timeout_seconds: float = shutdown_timeout_seconds - self._clock: Callable[[], datetime] = clock or (lambda: datetime.now(UTC)) - self._failure_callback: Callable[[BaseException], object] | None = failure_callback - self._queue: deque[CommandOutputRecord] = deque() - self._buffer_stream: CommandOutputStream | None = None - self._buffer_parts: list[str] = [] - self._buffer_bytes: int = 0 - self._buffer_started_at: float | None = None - self._partials: dict[CommandOutputStream, str] = { - CommandOutputStream.STDOUT: "", - CommandOutputStream.STDERR: "", - } - self._condition = threading.Condition() - self._sequence: int = 0 - self._accepted: int = 0 - self._delivered: int = 0 - self._failed: int = 0 - self._dropped: int = 0 - self._accepting: bool = True - self._stopping: bool = False - self._force_stop: bool = False - self._worker_finished: bool = False - self._suppressed_threads: set[int] = set() - self._thread = threading.Thread( - target=self._run, name="sqlbuild-output-exporter", daemon=True - ) - self._thread.start() - - def append(self, *, stream: CommandOutputStream, text: str) -> None: - """Accept one stream fragment without waiting for destination work.""" - - if not text or threading.get_ident() in self._suppressed_threads: - return - with self._condition: - if not self._accepting: - return - combined: str = self._partials[stream] + text - lines: list[str] = combined.splitlines(keepends=True) - trailing: str = "" - if lines and not lines[-1].endswith(("\n", "\r")): - trailing = lines.pop() - self._partials[stream] = trailing - for line in lines: - self._buffer_text_fragment( - stream=stream, text=line, immediate_priority=OutputRecordPriority.BULK - ) - self._condition.notify() - - def close(self) -> CommandOutputCaptureSummary: - """Flush partial records until a fixed deadline and return without indefinite waits.""" - - with self._condition: - if self._accepting: - for stream, partial in self._partials.items(): - if partial: - self._buffer_text_fragment( - stream=stream, - text=partial, - immediate_priority=OutputRecordPriority.TERMINAL, - ) - self._flush_buffer(priority=OutputRecordPriority.TERMINAL) - self._flush_buffer(priority=OutputRecordPriority.TERMINAL) - self._partials = { - CommandOutputStream.STDOUT: "", - CommandOutputStream.STDERR: "", - } - self._accepting = False - if self._dropped: - self._enqueue_loss_summary() - self._stopping = True - self._condition.notify_all() - deadline: float = time.monotonic() + self._shutdown_timeout_seconds - self._thread.join(timeout=max(0.0, deadline - time.monotonic())) - with self._condition: - if self._thread.is_alive(): - self._force_stop = True - self._dropped += len(self._queue) - self._queue.clear() - self._condition.notify_all() - return CommandOutputCaptureSummary( - accepted=self._accepted, - delivered=self._delivered, - failed=self._failed, - dropped=self._dropped, - queue_depth=len(self._queue), - queue_capacity=self._capacity, - flush_complete=self._worker_finished and not self._queue, - ) - - def _buffer_text_fragment( - self, - *, - stream: CommandOutputStream, - text: str, - immediate_priority: OutputRecordPriority, - ) -> None: - plain: str = strip_ansi(text) - chunks: tuple[str, ...] = chunk_text(text=plain, max_bytes=self._max_record_bytes) - if len(chunks) > 1: - self._flush_buffer(priority=OutputRecordPriority.BULK) - self._enqueue_chunks(stream=stream, chunks=chunks, priority=immediate_priority) - return - chunk: str = chunks[0] - if self._buffer_stream is not None and self._buffer_stream is not stream: - self._flush_buffer(priority=OutputRecordPriority.BULK) - chunk_bytes: int = len(chunk.encode("utf-8")) - if self._buffer_parts and self._buffer_bytes + chunk_bytes > self._max_record_bytes: - self._flush_buffer(priority=OutputRecordPriority.BULK) - if not self._buffer_parts: - self._buffer_stream = stream - self._buffer_started_at = time.monotonic() - self._buffer_parts.append(chunk) - self._buffer_bytes += chunk_bytes - - def _flush_buffer(self, *, priority: OutputRecordPriority) -> None: - if not self._buffer_parts or self._buffer_stream is None: - return - stream: CommandOutputStream = self._buffer_stream - text: str = "".join(self._buffer_parts) - self._buffer_stream = None - self._buffer_parts = [] - self._buffer_bytes = 0 - self._buffer_started_at = None - self._enqueue_chunks(stream=stream, chunks=(text,), priority=priority) - - def _enqueue_chunks( - self, - *, - stream: CommandOutputStream, - chunks: tuple[str, ...], - priority: OutputRecordPriority, - ) -> None: - for index, chunk in enumerate(chunks): - record: CommandOutputRecord = CommandOutputRecord( - invocation_id=self._invocation_id, - run_id=self._run_id, - sequence=self._next_sequence(), - occurred_at=self._clock(), - stream=stream, - message=chunk, - external_context=self._context, - chunk_index=index, - chunk_count=len(chunks), - priority=priority, - ) - self._put(record) - - def _put(self, record: CommandOutputRecord) -> None: - self._accepted += 1 - if len(self._queue) < self._capacity: - self._queue.append(record) - return - if record.priority is OutputRecordPriority.TERMINAL: - for queued in self._queue: - if queued.priority is OutputRecordPriority.BULK: - self._queue.remove(queued) - self._dropped += 1 - self._queue.append(record) - return - self._dropped += 1 - - def _enqueue_loss_summary(self) -> None: - summary_displaces_bulk: bool = len(self._queue) >= self._capacity and any( - queued.priority is OutputRecordPriority.BULK for queued in self._queue - ) - dropped_before_summary: int = self._dropped + int(summary_displaces_bulk) - record: CommandOutputRecord = CommandOutputRecord( - invocation_id=self._invocation_id, - run_id=self._run_id, - sequence=self._next_sequence(), - occurred_at=self._clock(), - stream=CommandOutputStream.STDERR, - message=f"SQLBuild output export dropped {dropped_before_summary} record(s)", - external_context=self._context, - priority=OutputRecordPriority.TERMINAL, - record_type=COMMAND_OUTPUT_LOSS_RECORD_TYPE, - dropped_records=dropped_before_summary, - ) - self._put(record) - - def _next_sequence(self) -> int: - sequence: int = self._sequence - self._sequence += 1 - return sequence - - def _run(self) -> None: - try: - while True: - with self._condition: - while not self._queue and not self._stopping: - if self._buffer_started_at is None: - self._condition.wait() - continue - remaining: float = ( - self._buffer_started_at - + self._flush_interval_seconds - - time.monotonic() - ) - if remaining > 0: - self._condition.wait(timeout=remaining) - continue - self._flush_buffer(priority=OutputRecordPriority.BULK) - if self._force_stop or (not self._queue and self._stopping): - return - batch: tuple[CommandOutputRecord, ...] = tuple( - self._queue.popleft() - for _ in range(min(self._batch_size, len(self._queue))) - ) - self._export(batch) - finally: - with self._condition: - self._worker_finished = True - self._condition.notify_all() - - def _export(self, batch: tuple[CommandOutputRecord, ...]) -> None: - thread_id: int = threading.get_ident() - self._suppressed_threads.add(thread_id) - try: - exporter: CommandOutputBatchExporter = cast(CommandOutputBatchExporter, self._exporter) - exporter.export_output(batch) - except BaseException as error: - with self._condition: - self._failed += len(batch) - if self._failure_callback is not None: - try: - self._failure_callback(error) - except BaseException: - pass - else: - with self._condition: - self._delivered += len(batch) - finally: - self._suppressed_threads.discard(thread_id) diff --git a/src/sqlbuild/runtime/output_capture/main/command_output_context.py b/src/sqlbuild/runtime/output_capture/main/command_output_context.py deleted file mode 100644 index 5ed24ff531..0000000000 --- a/src/sqlbuild/runtime/output_capture/main/command_output_context.py +++ /dev/null @@ -1,14 +0,0 @@ -"""Scope command-output integration context.""" - -from collections.abc import Iterator, Mapping -from contextlib import contextmanager - -from sqlbuild.runtime.output_capture._helpers.scope import output_capture_context - - -@contextmanager -def command_output_context(*, external_context: Mapping[str, object]) -> Iterator[None]: - """Attach opaque integration context to command-output records.""" - - with output_capture_context(external_context=external_context): - yield diff --git a/src/sqlbuild/runtime/output_capture/main/command_output_from_json.py b/src/sqlbuild/runtime/output_capture/main/command_output_from_json.py deleted file mode 100644 index 56f0aeeaeb..0000000000 --- a/src/sqlbuild/runtime/output_capture/main/command_output_from_json.py +++ /dev/null @@ -1,83 +0,0 @@ -"""Deserialize and validate one canonical command-output record.""" - -from __future__ import annotations - -import json -from datetime import datetime -from typing import Any, cast - -from sqlbuild.runtime.output_capture.constants import ( - COMMAND_OUTPUT_ENVELOPE_FIELDS, - COMMAND_OUTPUT_LOSS_RECORD_TYPE, - COMMAND_OUTPUT_RECORD_TYPE, - CURRENT_COMMAND_OUTPUT_SCHEMA_VERSION, -) -from sqlbuild.runtime.output_capture.exceptions import CommandOutputValidationError -from sqlbuild.runtime.output_capture.models import CommandOutputRecord -from sqlbuild.runtime.output_capture.types import CommandOutputStream, OutputRecordPriority - - -def command_output_from_json(raw_json: str) -> CommandOutputRecord: - """Deserialize one known command-output record strictly.""" - - try: - raw: object = json.loads(raw_json) - except (TypeError, ValueError) as error: - raise CommandOutputValidationError("command output must be valid JSON") from error - if not isinstance(raw, dict): - raise CommandOutputValidationError("command output must be a JSON object") - data: dict[str, Any] = cast(dict[str, Any], raw) - unknown: frozenset[str] = frozenset(data) - COMMAND_OUTPUT_ENVELOPE_FIELDS - missing: frozenset[str] = COMMAND_OUTPUT_ENVELOPE_FIELDS - frozenset(data) - if unknown or missing: - raise CommandOutputValidationError("command output fields do not match schema version 1") - if ( - isinstance(data["schema_version"], bool) - or not isinstance(data["schema_version"], int) - or data["schema_version"] != CURRENT_COMMAND_OUTPUT_SCHEMA_VERSION - ): - raise CommandOutputValidationError("unsupported command output schema version") - if not isinstance(data["record_type"], str) or data["record_type"] not in { - COMMAND_OUTPUT_RECORD_TYPE, - COMMAND_OUTPUT_LOSS_RECORD_TYPE, - }: - raise CommandOutputValidationError("unknown command output record type") - context: object = data["external_context"] - if not isinstance(context, dict): - raise CommandOutputValidationError("command output external_context must be an object") - try: - record: CommandOutputRecord = CommandOutputRecord( - invocation_id=data["invocation_id"], - run_id=data["run_id"], - sequence=data["sequence"], - occurred_at=_occurred_at(data["occurred_at"]), - stream=CommandOutputStream(data["stream"]), - message=data["message"], - external_context=context, - chunk_index=data["chunk_index"], - chunk_count=data["chunk_count"], - priority=OutputRecordPriority(data["priority"]), - record_type=data["record_type"], - dropped_records=data["dropped_records"], - schema_version=data["schema_version"], - producer=data["producer"], - producer_version=data["producer_version"], - ) - except (TypeError, ValueError) as error: - if isinstance(error, CommandOutputValidationError): - raise - raise CommandOutputValidationError( - "command output contains invalid field values" - ) from error - if data["record_id"] != record.record_id: - raise CommandOutputValidationError("command output record_id does not match its identity") - return record - - -def _occurred_at(value: object) -> datetime: - if not isinstance(value, str): - raise CommandOutputValidationError("occurred_at must be an ISO-8601 string") - try: - return datetime.fromisoformat(value.replace("Z", "+00:00")) - except ValueError as error: - raise CommandOutputValidationError("occurred_at must be an ISO-8601 string") from error diff --git a/src/sqlbuild/runtime/output_capture/main/command_output_to_json.py b/src/sqlbuild/runtime/output_capture/main/command_output_to_json.py deleted file mode 100644 index 79c8039649..0000000000 --- a/src/sqlbuild/runtime/output_capture/main/command_output_to_json.py +++ /dev/null @@ -1,47 +0,0 @@ -"""Serialize one command-output record into its canonical wire format.""" - -from __future__ import annotations - -import json -from collections.abc import Mapping - -from sqlbuild.runtime.output_capture.models import CommandOutputRecord - - -def _json_data(value: object) -> object: - if isinstance(value, Mapping): - return {key: _json_data(item) for key, item in value.items()} - if isinstance(value, tuple): - return [_json_data(item) for item in value] - return value - - -def command_output_to_json(record: CommandOutputRecord) -> str: - """Serialize one command-output record deterministically.""" - - return json.dumps( - { - "record_id": record.record_id, - "record_type": record.record_type, - "schema_version": record.schema_version, - "producer": record.producer, - "producer_version": record.producer_version, - "occurred_at": record.occurred_at.isoformat(timespec="microseconds").replace( - "+00:00", "Z" - ), - "invocation_id": record.invocation_id, - "run_id": record.run_id, - "sequence": record.sequence, - "stream": record.stream.value, - "message": record.message, - "chunk_index": record.chunk_index, - "chunk_count": record.chunk_count, - "priority": record.priority.value, - "dropped_records": record.dropped_records, - "external_context": _json_data(record.external_context), - }, - ensure_ascii=True, - allow_nan=False, - separators=(",", ":"), - sort_keys=True, - ) diff --git a/src/sqlbuild/cli/commands/main/entrypoint/_dispatch_with_output_capture.py b/src/sqlbuild/runtime/output_capture/main/configured_output_capture_scope.py similarity index 100% rename from src/sqlbuild/cli/commands/main/entrypoint/_dispatch_with_output_capture.py rename to src/sqlbuild/runtime/output_capture/main/configured_output_capture_scope.py diff --git a/src/sqlbuild/runtime/output_capture/main/current_output_capture_context.py b/src/sqlbuild/runtime/output_capture/main/current_output_capture_context.py deleted file mode 100644 index 809f4a1ffb..0000000000 --- a/src/sqlbuild/runtime/output_capture/main/current_output_capture_context.py +++ /dev/null @@ -1,12 +0,0 @@ -"""Active output capture context lookup.""" - -from sqlbuild.runtime.output_capture._helpers.scope import ( - current_output_capture_context as _current_output_capture_context, -) -from sqlbuild.runtime.output_capture.models import OutputCaptureContext - - -def current_output_capture_context() -> OutputCaptureContext | None: - """Return optional integration context for the active command.""" - - return _current_output_capture_context() diff --git a/src/sqlbuild/runtime/output_capture/main/invocation_context.py b/src/sqlbuild/runtime/output_capture/main/invocation_context.py deleted file mode 100644 index 6fac667690..0000000000 --- a/src/sqlbuild/runtime/output_capture/main/invocation_context.py +++ /dev/null @@ -1,61 +0,0 @@ -"""Load bounded opaque invocation context for subprocess integrations.""" - -from __future__ import annotations - -import json -import os -from collections.abc import Mapping -from typing import cast - -from sqlbuild.runtime.output_capture._helpers.json import freeze_command_output_json -from sqlbuild.runtime.output_capture.constants import ( - INVOCATION_CONTEXT_ENV, - MAX_INVOCATION_CONTEXT_BYTES, - MAX_INVOCATION_CONTEXT_DEPTH, -) -from sqlbuild.runtime.output_capture.exceptions import OutputCaptureInputError - - -def invocation_context_from_environment( - *, environment: Mapping[str, str] | None = None -) -> Mapping[str, object]: - """Return validated opaque invocation metadata supplied by an integration.""" - - raw: str | None = (os.environ if environment is None else environment).get( - INVOCATION_CONTEXT_ENV - ) - if raw is None: - return {} - if len(raw.encode("utf-8")) > MAX_INVOCATION_CONTEXT_BYTES: - raise OutputCaptureInputError( - f"{INVOCATION_CONTEXT_ENV} must not exceed {MAX_INVOCATION_CONTEXT_BYTES} bytes" - ) - try: - loaded: object = json.loads(raw) - except json.JSONDecodeError as error: - raise OutputCaptureInputError( - f"{INVOCATION_CONTEXT_ENV} must contain valid JSON" - ) from error - if not isinstance(loaded, dict): - raise OutputCaptureInputError(f"{INVOCATION_CONTEXT_ENV} must contain a JSON object") - _validate_depth(value=loaded, depth=1) - try: - frozen: object = freeze_command_output_json(value=loaded, path="external_context") - except ValueError as error: - raise OutputCaptureInputError(str(error)) from error - if not isinstance(frozen, Mapping): - raise OutputCaptureInputError(f"{INVOCATION_CONTEXT_ENV} must contain a JSON object") - return cast(Mapping[str, object], frozen) - - -def _validate_depth(*, value: object, depth: int) -> None: - if depth > MAX_INVOCATION_CONTEXT_DEPTH: - raise OutputCaptureInputError( - f"{INVOCATION_CONTEXT_ENV} must not exceed {MAX_INVOCATION_CONTEXT_DEPTH} levels" - ) - if isinstance(value, dict): - for item in value.values(): - _validate_depth(value=item, depth=depth + 1) - elif isinstance(value, list): - for item in value: - _validate_depth(value=item, depth=depth + 1) From 7372457797587bf542c25203eed7c09ebc3cb0b9 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 15 Oct 2026 10:55:52 +0100 Subject: [PATCH 10/12] refactor: host command-output capture natively Move command-output text normalization, chunking, line buffering, the bounded priority queue, record validation, wire JSON, the invocation context and the delivery dispatcher into Rust. The dispatcher keeps Python threading primitives so delivery still runs on a named Python daemon thread. Python keeps the stream tees, the scope that installs them, and the public record shapes. --- .../src/command_output/_helpers/buffering.rs | 110 ++++ .../src/command_output/_helpers/mod.rs | 2 + .../src/command_output/_helpers/text.rs | 235 ++++++++ .../src/command_output/constants.rs | 36 ++ .../src/command_output/errors.rs | 10 + .../src/command_output/main/append_output.rs | 33 ++ .../command_output/main/chunk_output_text.rs | 10 + .../src/command_output/main/close_output.rs | 26 + .../command_output/main/flush_idle_output.rs | 18 + .../src/command_output/main/mod.rs | 7 + .../main/output_loss_summary.rs | 30 + .../main/queue_output_record.rs | 29 + .../main/validate_output_record.rs | 61 +++ .../sqlbuild-core/src/command_output/mod.rs | 11 + .../src/command_output/models.rs | 125 +++++ .../sqlbuild-core/src/command_output/tests.rs | 6 + .../command_output/tests/test_buffering.rs | 63 +++ .../src/command_output/tests/test_text.rs | 64 +++ .../src/command_output/tests/test_types.rs | 15 + .../sqlbuild-core/src/command_output/types.rs | 4 + crates/sqlbuild-core/src/lib.rs | 1 + .../src/bindings/_helpers/cli/entry.rs | 10 +- .../bindings/_helpers/cli/observability.rs | 2 +- .../src/bindings/_helpers/cli/registration.rs | 5 +- .../src/bindings/_helpers/mod.rs | 1 + .../_helpers/output_capture/context.rs | 88 +++ .../_helpers/output_capture/dispatcher.rs | 512 ++++++++++++++++++ .../bindings/_helpers/output_capture/mod.rs | 6 + .../_helpers/output_capture/records.rs | 372 +++++++++++++ .../_helpers/output_capture/registration.rs | 12 + .../bindings/_helpers/output_capture/text.rs | 31 ++ .../_helpers/output_capture/values.rs | 178 ++++++ python-boundary.toml | 41 +- src/sqlbuild/_native.pyi | 43 ++ .../output_capture/classes/binary_tee.py | 2 +- .../output_capture/classes/text_tee.py | 2 +- .../main/configured_output_capture_scope.py | 12 +- src/sqlbuild/runtime/output_capture/models.py | 71 +-- src/sqlbuild/sinks.py | 12 +- .../test_dispatch_with_output_capture.py | 6 +- .../output_capture/_helpers/test_text.py | 2 +- .../runtime/output_capture/classes/helpers.py | 2 +- .../classes/test_output_capture_dispatcher.py | 2 +- .../main/test_invocation_context.py | 4 +- 44 files changed, 2198 insertions(+), 114 deletions(-) create mode 100644 crates/sqlbuild-core/src/command_output/_helpers/buffering.rs create mode 100644 crates/sqlbuild-core/src/command_output/_helpers/mod.rs create mode 100644 crates/sqlbuild-core/src/command_output/_helpers/text.rs create mode 100644 crates/sqlbuild-core/src/command_output/constants.rs create mode 100644 crates/sqlbuild-core/src/command_output/errors.rs create mode 100644 crates/sqlbuild-core/src/command_output/main/append_output.rs create mode 100644 crates/sqlbuild-core/src/command_output/main/chunk_output_text.rs create mode 100644 crates/sqlbuild-core/src/command_output/main/close_output.rs create mode 100644 crates/sqlbuild-core/src/command_output/main/flush_idle_output.rs create mode 100644 crates/sqlbuild-core/src/command_output/main/mod.rs create mode 100644 crates/sqlbuild-core/src/command_output/main/output_loss_summary.rs create mode 100644 crates/sqlbuild-core/src/command_output/main/queue_output_record.rs create mode 100644 crates/sqlbuild-core/src/command_output/main/validate_output_record.rs create mode 100644 crates/sqlbuild-core/src/command_output/mod.rs create mode 100644 crates/sqlbuild-core/src/command_output/models.rs create mode 100644 crates/sqlbuild-core/src/command_output/tests.rs create mode 100644 crates/sqlbuild-core/src/command_output/tests/test_buffering.rs create mode 100644 crates/sqlbuild-core/src/command_output/tests/test_text.rs create mode 100644 crates/sqlbuild-core/src/command_output/tests/test_types.rs create mode 100644 crates/sqlbuild-core/src/command_output/types.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/output_capture/context.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/output_capture/dispatcher.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/output_capture/mod.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/output_capture/records.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/output_capture/registration.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/output_capture/text.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/output_capture/values.rs diff --git a/crates/sqlbuild-core/src/command_output/_helpers/buffering.rs b/crates/sqlbuild-core/src/command_output/_helpers/buffering.rs new file mode 100644 index 0000000000..1d13a20d1e --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/_helpers/buffering.rs @@ -0,0 +1,110 @@ +//! `OutputCaptureDispatcher` fragment buffering: chunk, coalesce and flush into record drafts. + +use crate::command_output::_helpers::text::{chunk_text, strict_utf8_length, strip_ansi}; +use crate::command_output::errors::UnencodableText; +use crate::command_output::models::{ + BufferOutcome, OutputBuffer, OutputStream, RecordDraft, RecordPriority, +}; +use crate::command_output::types::PassText; + +/// `_enqueue_chunks`: one draft per chunk with consecutive sequence numbers. +fn enqueue_chunks( + mut outcome: BufferOutcome, + (stream, chunks): (OutputStream, Vec), + priority: RecordPriority, +) -> BufferOutcome { + let chunk_count = chunks.len(); + for (chunk_index, message) in chunks.into_iter().enumerate() { + let sequence = outcome.buffer.sequence; + outcome.buffer.sequence += 1; + outcome.drafts.push(RecordDraft { + sequence, + stream, + message, + chunk_index, + chunk_count, + priority, + dropped_records: 0, + }); + } + outcome +} + +/// `_flush_buffer`: turn the coalesced buffer into one record. +pub(crate) fn flush_buffer(mut outcome: BufferOutcome, priority: RecordPriority) -> BufferOutcome { + let Some(stream) = outcome.buffer.buffer_stream else { + return outcome; + }; + if outcome.buffer.buffer_part_count == 0 { + return outcome; + } + let text = std::mem::take(&mut outcome.buffer.buffer_text); + outcome.buffer.buffer_stream = None; + outcome.buffer.buffer_part_count = 0; + outcome.buffer.buffer_bytes = 0; + outcome.buffer.buffer_started_at = None; + enqueue_chunks(outcome, (stream, vec![text]), priority) +} + +fn failed(mut outcome: BufferOutcome, error: UnencodableText) -> BufferOutcome { + outcome.error = Some(error); + outcome +} + +/// `_buffer_text_fragment`: chunk and coalesce one ANSI-free fragment; failures end the outcome. +pub(crate) fn buffer_fragment( + mut outcome: BufferOutcome, + (stream, text): (OutputStream, &[u8]), + (immediate, now): (RecordPriority, f64), +) -> BufferOutcome { + let mut chunks = match chunk_text(&strip_ansi(text), outcome.buffer.max_record_bytes) { + Ok(chunks) => chunks, + Err(error) => return failed(outcome, error), + }; + if chunks.len() > 1 { + outcome = flush_buffer(outcome, RecordPriority::Bulk); + return enqueue_chunks(outcome, (stream, chunks), immediate); + } + let chunk = chunks.pop().unwrap_or_default(); + if outcome + .buffer + .buffer_stream + .is_some_and(|buffered| buffered != stream) + { + outcome = flush_buffer(outcome, RecordPriority::Bulk); + } + let chunk_bytes = match strict_utf8_length(&chunk) { + Ok(length) => length, + Err(error) => return failed(outcome, error), + }; + let exceeds = i64::try_from(outcome.buffer.buffer_bytes + chunk_bytes).unwrap_or(i64::MAX) + > outcome.buffer.max_record_bytes; + if outcome.buffer.buffer_part_count > 0 && exceeds { + outcome = flush_buffer(outcome, RecordPriority::Bulk); + } + if outcome.buffer.buffer_part_count == 0 { + outcome.buffer.buffer_stream = Some(stream); + outcome.buffer.buffer_started_at = Some(now); + } + outcome.buffer.buffer_text.extend_from_slice(&chunk); + outcome.buffer.buffer_part_count += 1; + outcome.buffer.buffer_bytes += chunk_bytes; + outcome +} + +/// A fresh outcome continuing from `buffer`. +pub(crate) fn start(buffer: OutputBuffer) -> BufferOutcome { + BufferOutcome { + buffer, + drafts: Vec::new(), + error: None, + } +} + +/// The index of `stream` in per-stream arrays: stdout, then stderr. +pub(crate) fn stream_slot(stream: OutputStream) -> usize { + match stream { + OutputStream::Stdout => 0, + OutputStream::Stderr => 1, + } +} diff --git a/crates/sqlbuild-core/src/command_output/_helpers/mod.rs b/crates/sqlbuild-core/src/command_output/_helpers/mod.rs new file mode 100644 index 0000000000..cd9ceb44bd --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/_helpers/mod.rs @@ -0,0 +1,2 @@ +pub(crate) mod buffering; +pub(crate) mod text; diff --git a/crates/sqlbuild-core/src/command_output/_helpers/text.rs b/crates/sqlbuild-core/src/command_output/_helpers/text.rs new file mode 100644 index 0000000000..4a01546a74 --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/_helpers/text.rs @@ -0,0 +1,235 @@ +//! `strip_ansi`, `str.splitlines(keepends=True)` and `chunk_text` over `surrogatepass` bytes. + +use crate::command_output::constants::{ + ASCII_LINE_BREAKS, BELL, CSI_FINAL, CSI_INTERMEDIATES, CSI_INTRODUCER, CSI_PARAMETERS, + ESCAPABLE_SURROGATES, ESCAPE, NEXT_LINE, OSC_INTRODUCER, STRING_TERMINATOR, SURROGATE_LEAD, + SURROGATE_SECOND, UNICODE_LINE_BREAKS, UTF8_CONTINUATION_MASK, UTF8_CONTINUATION_TAG, +}; +use crate::command_output::errors::UnencodableText; + +/// Byte length of the code point whose UTF-8 encoding starts with `lead`. +fn sequence_length(lead: u8) -> usize { + match lead { + 0x00..=0x7f => 1, + 0xc0..=0xdf => 2, + 0xe0..=0xef => 3, + _ => 4, + } +} + +fn is_continuation(byte: u8) -> bool { + byte & UTF8_CONTINUATION_MASK == UTF8_CONTINUATION_TAG +} + +/// The surrogate code point encoded at `index`, if the bytes there encode one. +fn surrogate_at(text: &[u8], index: usize) -> Option { + let (Some(&lead), Some(&second), Some(&third)) = + (text.get(index), text.get(index + 1), text.get(index + 2)) + else { + return None; + }; + if lead != SURROGATE_LEAD || !SURROGATE_SECOND.contains(&second) { + return None; + } + Some(0xd000 | (u32::from(second & 0x3f) << 6) | u32::from(third & 0x3f)) +} + +/// The first lone surrogate in `text`, which strict UTF-8 cannot encode. +pub(crate) fn first_surrogate(text: &[u8]) -> Option { + let found = text + .windows(2) + .position(|pair| pair[0] == SURROGATE_LEAD && SURROGATE_SECOND.contains(&pair[1]))?; + let code_point = surrogate_at(text, found)?; + let position = text[..found] + .iter() + .filter(|byte| !is_continuation(**byte)) + .count(); + Some(UnencodableText { + position, + code_point, + }) +} + +/// `len(text.encode("utf-8"))`, failing on a lone surrogate. +pub(crate) fn strict_utf8_length(text: &[u8]) -> Result { + match first_surrogate(text) { + Some(error) => Err(error), + None => Ok(text.len()), + } +} + +/// End of an OSC sequence starting at `start`, terminated by BEL or `ESC \`. +fn osc_end(text: &[u8], start: usize) -> Option { + let mut index = start + 2; + while let Some(&byte) = text.get(index) { + if byte == BELL { + return Some(index + 1); + } + if byte == ESCAPE { + return (text.get(index + 1) == Some(&STRING_TERMINATOR)).then_some(index + 2); + } + index += 1; + } + None +} + +/// End of a CSI sequence starting at `start`: parameters, intermediates, one final byte. +fn csi_end(text: &[u8], start: usize) -> Option { + let mut index = start + 2; + while text + .get(index) + .is_some_and(|byte| CSI_PARAMETERS.contains(byte)) + { + index += 1; + } + while text + .get(index) + .is_some_and(|byte| CSI_INTERMEDIATES.contains(byte)) + { + index += 1; + } + text.get(index) + .filter(|byte| CSI_FINAL.contains(byte)) + .map(|_| index + 1) +} + +/// `_ANSI_ESCAPE.sub("", text)`: drop OSC and CSI terminal control sequences. +pub(crate) fn strip_ansi(text: &[u8]) -> Vec { + let mut plain: Vec = Vec::with_capacity(text.len()); + let mut index: usize = 0; + while index < text.len() { + let end = match (text[index], text.get(index + 1)) { + (ESCAPE, Some(&OSC_INTRODUCER)) => osc_end(text, index), + (ESCAPE, Some(&CSI_INTRODUCER)) => csi_end(text, index), + _ => None, + }; + match end { + Some(end) => index = end, + None => { + plain.push(text[index]); + index += 1; + } + } + } + plain +} + +/// Length of the line boundary at `index`, if one starts there. +fn line_break_length(text: &[u8], index: usize) -> Option { + let byte = text[index]; + if byte == b'\r' && text.get(index + 1) == Some(&b'\n') { + return Some(2); + } + if ASCII_LINE_BREAKS.contains(&byte) { + return Some(1); + } + if text[index..].starts_with(&NEXT_LINE) { + return Some(NEXT_LINE.len()); + } + UNICODE_LINE_BREAKS + .iter() + .find(|boundary| text[index..].starts_with(boundary.as_slice())) + .map(|boundary| boundary.len()) +} + +/// `text.splitlines(keepends=True)`. +pub(crate) fn split_lines(text: &[u8]) -> Vec<&[u8]> { + let mut lines: Vec<&[u8]> = Vec::new(); + let mut start: usize = 0; + let mut index: usize = 0; + while index < text.len() { + match line_break_length(text, index) { + Some(length) => { + index += length; + lines.push(&text[start..index]); + start = index; + } + None => index += sequence_length(text[index]), + } + } + if start < text.len() { + lines.push(&text[start..]); + } + lines +} + +/// `_chunk_utf8`: greedy chunks ending at the last code point boundary within `max_bytes`. +fn chunk_utf8(text: &[u8], max_bytes: i64) -> Vec> { + let mut chunks: Vec> = Vec::new(); + let total = text.len(); + let mut start: usize = 0; + while start < total { + let limit = usize::try_from(max_bytes).unwrap_or(0); + let mut end = start.saturating_add(limit); + if end >= total { + end = total; + } else { + while end > start && is_continuation(text[end]) { + end -= 1; + } + if end <= start { + end = start + 1; + while end < total && is_continuation(text[end]) { + end += 1; + } + } + } + chunks.push(text[start..end].to_vec()); + start = end; + } + if chunks.is_empty() { + chunks.push(Vec::new()); + } + chunks +} + +/// `len(character.encode("utf-8", "surrogateescape"))` of the code point at `index`. +fn escaped_length(text: &[u8], index: usize, position: usize) -> Result { + match surrogate_at(text, index) { + Some(code_point) if ESCAPABLE_SURROGATES.contains(&code_point) => Ok(1), + Some(code_point) => Err(UnencodableText { + position, + code_point, + }), + None => Ok(sequence_length(text[index])), + } +} + +/// `_chunk_text_by_character`: text with surrogates, each escaped surrogate one byte. +fn chunk_by_character(text: &[u8], max_bytes: i64) -> Result>, UnencodableText> { + let mut chunks: Vec> = Vec::new(); + let mut current: Vec = Vec::new(); + let mut current_bytes: i64 = 0; + let mut index: usize = 0; + let mut position: usize = 0; + while index < text.len() { + let length = sequence_length(text[index]); + let encoded = i64::try_from(escaped_length(text, index, position)?).unwrap_or(i64::MAX); + if !current.is_empty() && current_bytes + encoded > max_bytes { + chunks.push(std::mem::take(&mut current)); + current_bytes = 0; + } + current.extend_from_slice(&text[index..index + length]); + current_bytes += encoded; + index += length; + position += 1; + } + if !current.is_empty() || chunks.is_empty() { + chunks.push(current); + } + Ok(chunks) +} + +/// `chunk_text`: split text deterministically without splitting code points. +pub(crate) fn chunk_text(text: &[u8], max_bytes: i64) -> Result>, UnencodableText> { + if text.is_ascii() && max_bytes > 0 { + return Ok(chunk_utf8(text, max_bytes)); + } + if first_surrogate(text).is_some() { + return chunk_by_character(text, max_bytes); + } + if i64::try_from(text.len()).unwrap_or(i64::MAX) <= max_bytes { + return Ok(vec![text.to_vec()]); + } + Ok(chunk_utf8(text, max_bytes)) +} diff --git a/crates/sqlbuild-core/src/command_output/constants.rs b/crates/sqlbuild-core/src/command_output/constants.rs new file mode 100644 index 0000000000..627d6f538e --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/constants.rs @@ -0,0 +1,36 @@ +//! Python's `sqlbuild.runtime.output_capture.constants` delivery limits and text bytes. + +pub const MIN_OUTPUT_RECORD_BYTES: i64 = 4; +pub const FLUSH_INTERVAL_SECONDS: f64 = 0.25; +/// Seconds left before a buffered record is due; at or below this it is flushed. +pub const NO_SECONDS_REMAINING: f64 = 0.0; +/// The shortest shutdown wait a dispatcher accepts. +pub const MIN_SHUTDOWN_TIMEOUT_SECONDS: f64 = 0.0; +pub const COMMAND_OUTPUT_RECORD_TYPE: &str = "command_output"; +pub const COMMAND_OUTPUT_LOSS_RECORD_TYPE: &str = "command_output_loss"; +pub const CURRENT_SCHEMA_VERSION: i64 = 1; +pub const MAX_INVOCATION_CONTEXT_BYTES: usize = 16 * 1024; +pub const MAX_INVOCATION_CONTEXT_DEPTH: usize = 8; +pub const INVOCATION_CONTEXT_ENV: &str = "SQLBUILD_INVOCATION_CONTEXT_JSON"; +pub(crate) const ESCAPE: u8 = 0x1b; +pub(crate) const BELL: u8 = 0x07; +pub(crate) const OSC_INTRODUCER: u8 = b']'; +pub(crate) const CSI_INTRODUCER: u8 = b'['; +pub(crate) const STRING_TERMINATOR: u8 = b'\\'; +pub(crate) const CSI_PARAMETERS: std::ops::RangeInclusive = 0x30..=0x3f; +pub(crate) const CSI_INTERMEDIATES: std::ops::RangeInclusive = 0x20..=0x2f; +pub(crate) const CSI_FINAL: std::ops::RangeInclusive = 0x40..=0x7e; +/// Single-byte line boundaries of `str.splitlines`: `\n \v \f \r \x1c \x1d \x1e`. +pub(crate) const ASCII_LINE_BREAKS: [u8; 7] = [0x0a, 0x0b, 0x0c, 0x0d, 0x1c, 0x1d, 0x1e]; +/// UTF-8 of `\x85`. +pub(crate) const NEXT_LINE: [u8; 2] = [0xc2, 0x85]; +/// UTF-8 of `\u2028` and `\u2029`. +pub(crate) const UNICODE_LINE_BREAKS: [[u8; 3]; 2] = [[0xe2, 0x80, 0xa8], [0xe2, 0x80, 0xa9]]; +pub(crate) const UTF8_CONTINUATION_MASK: u8 = 0xc0; +pub(crate) const UTF8_CONTINUATION_TAG: u8 = 0x80; +/// Lead byte of a UTF-8 surrogate encoding (`surrogatepass`). +pub(crate) const SURROGATE_LEAD: u8 = 0xed; +/// Second bytes `0xa0..=0xbf` after `0xed` encode surrogates. +pub(crate) const SURROGATE_SECOND: std::ops::RangeInclusive = 0xa0..=0xbf; +/// `surrogateescape` maps exactly the low surrogates `U+DC80..=U+DCFF` to single bytes. +pub(crate) const ESCAPABLE_SURROGATES: std::ops::RangeInclusive = 0xdc80..=0xdcff; diff --git a/crates/sqlbuild-core/src/command_output/errors.rs b/crates/sqlbuild-core/src/command_output/errors.rs new file mode 100644 index 0000000000..cd6af8e579 --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/errors.rs @@ -0,0 +1,10 @@ +//! Text that cannot be encoded where Python's capture encodes it. + +/// Python's `UnicodeEncodeError` for text holding a surrogate the capture must encode. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct UnencodableText { + /// Code point index of the first surrogate that failed to encode. + pub position: usize, + /// The surrogate's code point. + pub code_point: u32, +} diff --git a/crates/sqlbuild-core/src/command_output/main/append_output.rs b/crates/sqlbuild-core/src/command_output/main/append_output.rs new file mode 100644 index 0000000000..3e6e89f340 --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/main/append_output.rs @@ -0,0 +1,33 @@ +//! `OutputCaptureDispatcher.append`: split a stream fragment into complete lines and buffer them. + +use crate::command_output::_helpers::buffering::{buffer_fragment, start, stream_slot}; +use crate::command_output::_helpers::text::split_lines; +use crate::command_output::models::{BufferOutcome, OutputBuffer, OutputStream, RecordPriority}; + +/// Buffer the complete lines of `text`, keeping an unterminated tail per stream. +pub fn append_output( + mut buffer: OutputBuffer, + (stream, text): (OutputStream, &[u8]), + now: f64, +) -> BufferOutcome { + let slot = stream_slot(stream); + let mut combined = std::mem::take(&mut buffer.partials[slot]); + combined.extend_from_slice(text); + let mut lines = split_lines(&combined); + let complete = lines + .last() + .is_some_and(|line| line.ends_with(b"\n") || line.ends_with(b"\r")); + let trailing: Vec = match complete { + true => Vec::new(), + false => lines.pop().map(<[u8]>::to_vec).unwrap_or_default(), + }; + buffer.partials[slot] = trailing; + let mut outcome = start(buffer); + for line in lines { + outcome = buffer_fragment(outcome, (stream, line), (RecordPriority::Bulk, now)); + if outcome.error.is_some() { + break; + } + } + outcome +} diff --git a/crates/sqlbuild-core/src/command_output/main/chunk_output_text.rs b/crates/sqlbuild-core/src/command_output/main/chunk_output_text.rs new file mode 100644 index 0000000000..3c0e565b9c --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/main/chunk_output_text.rs @@ -0,0 +1,10 @@ +//! `chunk_text`: split captured text deterministically without splitting code points. + +use crate::command_output::_helpers::text::chunk_text; +use crate::command_output::errors::UnencodableText; +use crate::command_output::types::PassText; + +/// Chunks of at most `max_bytes` UTF-8 bytes; a wider code point is a chunk of its own. +pub fn chunk_output_text(text: &[u8], max_bytes: i64) -> Result, UnencodableText> { + chunk_text(text, max_bytes) +} diff --git a/crates/sqlbuild-core/src/command_output/main/close_output.rs b/crates/sqlbuild-core/src/command_output/main/close_output.rs new file mode 100644 index 0000000000..9407d8f09d --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/main/close_output.rs @@ -0,0 +1,26 @@ +//! `OutputCaptureDispatcher.close`: flush unterminated tails and the buffer as terminal records. + +use crate::command_output::_helpers::buffering::{buffer_fragment, flush_buffer, start}; +use crate::command_output::models::{BufferOutcome, OutputBuffer, OutputStream, RecordPriority}; + +/// Flush each stream's partial line and the coalesced buffer with terminal priority. +pub fn close_output(buffer: OutputBuffer, now: f64) -> BufferOutcome { + let partials = buffer.partials.clone(); + let mut outcome = start(buffer); + for (stream, partial) in [OutputStream::Stdout, OutputStream::Stderr] + .into_iter() + .zip(partials) + { + if partial.is_empty() { + continue; + } + outcome = buffer_fragment(outcome, (stream, &partial), (RecordPriority::Terminal, now)); + if outcome.error.is_some() { + return outcome; + } + outcome = flush_buffer(outcome, RecordPriority::Terminal); + } + outcome = flush_buffer(outcome, RecordPriority::Terminal); + outcome.buffer.partials = [Vec::new(), Vec::new()]; + outcome +} diff --git a/crates/sqlbuild-core/src/command_output/main/flush_idle_output.rs b/crates/sqlbuild-core/src/command_output/main/flush_idle_output.rs new file mode 100644 index 0000000000..1833109c66 --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/main/flush_idle_output.rs @@ -0,0 +1,18 @@ +//! The delivery worker's idle step: wait for records or flush a buffered record that is due. + +use crate::command_output::_helpers::buffering::{flush_buffer, start}; +use crate::command_output::constants::{FLUSH_INTERVAL_SECONDS, NO_SECONDS_REMAINING}; +use crate::command_output::models::{IdleStep, OutputBuffer, RecordPriority}; + +/// Wait while nothing is buffered or the buffer is younger than the flush interval. +pub fn flush_idle_output(buffer: OutputBuffer, now: f64) -> (OutputBuffer, IdleStep) { + let Some(started_at) = buffer.buffer_started_at else { + return (buffer, IdleStep::Wait(None)); + }; + let remaining = started_at + FLUSH_INTERVAL_SECONDS - now; + if remaining > NO_SECONDS_REMAINING { + return (buffer, IdleStep::Wait(Some(remaining))); + } + let outcome = flush_buffer(start(buffer), RecordPriority::Bulk); + (outcome.buffer, IdleStep::Flushed(outcome.drafts)) +} diff --git a/crates/sqlbuild-core/src/command_output/main/mod.rs b/crates/sqlbuild-core/src/command_output/main/mod.rs new file mode 100644 index 0000000000..9b7cef61f9 --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/main/mod.rs @@ -0,0 +1,7 @@ +pub mod append_output; +pub mod chunk_output_text; +pub mod close_output; +pub mod flush_idle_output; +pub mod output_loss_summary; +pub mod queue_output_record; +pub mod validate_output_record; diff --git a/crates/sqlbuild-core/src/command_output/main/output_loss_summary.rs b/crates/sqlbuild-core/src/command_output/main/output_loss_summary.rs new file mode 100644 index 0000000000..9f30269a98 --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/main/output_loss_summary.rs @@ -0,0 +1,30 @@ +//! `_enqueue_loss_summary`: the terminal stderr record reporting dropped output. + +use crate::command_output::models::{ + OutputBuffer, OutputQueue, OutputStream, RecordDraft, RecordPriority, +}; + +/// The loss record counting drops so far, including the bulk record it is about to displace. +pub fn output_loss_summary( + mut buffer: OutputBuffer, + queue: &OutputQueue, +) -> (OutputBuffer, RecordDraft) { + let displaces_bulk = queue.items.len() >= queue.capacity + && queue + .items + .iter() + .any(|(_, priority)| *priority == RecordPriority::Bulk); + let dropped = queue.dropped + u64::from(displaces_bulk); + let sequence = buffer.sequence; + buffer.sequence += 1; + let draft = RecordDraft { + sequence, + stream: OutputStream::Stderr, + message: format!("SQLBuild output export dropped {dropped} record(s)").into_bytes(), + chunk_index: 0, + chunk_count: 1, + priority: RecordPriority::Terminal, + dropped_records: dropped, + }; + (buffer, draft) +} diff --git a/crates/sqlbuild-core/src/command_output/main/queue_output_record.rs b/crates/sqlbuild-core/src/command_output/main/queue_output_record.rs new file mode 100644 index 0000000000..d6ffd92ec6 --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/main/queue_output_record.rs @@ -0,0 +1,29 @@ +//! `_put`: accept one record into the bounded queue, displacing bulk output for terminal records. + +use crate::command_output::models::{OutputQueue, RecordPriority}; + +/// Queue `item`, or drop it (or the oldest bulk record for a terminal one) when full. +pub fn queue_output_record( + mut queue: OutputQueue, + item: T, + priority: RecordPriority, +) -> OutputQueue { + queue.accepted += 1; + if queue.items.len() < queue.capacity { + queue.items.push_back((item, priority)); + return queue; + } + if priority == RecordPriority::Terminal + && let Some(position) = queue + .items + .iter() + .position(|(_, queued)| *queued == RecordPriority::Bulk) + { + queue.items.remove(position); + queue.dropped += 1; + queue.items.push_back((item, priority)); + return queue; + } + queue.dropped += 1; + queue +} diff --git a/crates/sqlbuild-core/src/command_output/main/validate_output_record.rs b/crates/sqlbuild-core/src/command_output/main/validate_output_record.rs new file mode 100644 index 0000000000..5cb51ba1a1 --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/main/validate_output_record.rs @@ -0,0 +1,61 @@ +//! `CommandOutputRecord.__post_init__`: the command-output wire contract, field by field. + +use crate::command_output::constants::{ + COMMAND_OUTPUT_LOSS_RECORD_TYPE, COMMAND_OUTPUT_RECORD_TYPE, CURRENT_SCHEMA_VERSION, +}; +use crate::command_output::models::RecordShape; + +fn at_least(name: &str, value: Option, minimum: i64) -> Result { + value + .filter(|value| *value >= minimum) + .ok_or_else(|| format!("{name} must be an integer of at least {minimum}")) +} + +/// The first contract violation, in the order Python checks fields. +pub fn validate_output_record(shape: &RecordShape) -> Result<(), String> { + for (name, value) in &shape.identity { + if value.as_deref().is_none_or(str::is_empty) { + return Err(format!("{name} must be a non-empty string")); + } + } + if let Some(run_id) = &shape.run_id + && run_id.as_deref().is_none_or(str::is_empty) + { + return Err("run_id must be null or a non-empty string".to_owned()); + } + if shape.sequence.is_none_or(|sequence| sequence < 0) { + return Err("sequence must be a non-negative integer".to_owned()); + } + if !shape.occurred_at_is_utc { + return Err("occurred_at must be a timezone-aware UTC datetime".to_owned()); + } + if !shape.stream_is_known { + return Err("stream must be stdout or stderr".to_owned()); + } + if !shape.message_is_text { + return Err("message must be a string".to_owned()); + } + let chunk_index = at_least("chunk_index", shape.chunk_index, 0)?; + let chunk_count = at_least("chunk_count", shape.chunk_count, 1)?; + let dropped = at_least("dropped_records", shape.dropped_records, 0)?; + if chunk_index >= chunk_count { + return Err("chunk_index must be less than chunk_count".to_owned()); + } + if !shape.priority_is_known { + return Err("priority must be bulk or terminal".to_owned()); + } + match shape.record_type.as_deref() { + Some(COMMAND_OUTPUT_RECORD_TYPE) if dropped != 0 => { + return Err("command_output records must have zero dropped_records".to_owned()); + } + Some(COMMAND_OUTPUT_LOSS_RECORD_TYPE) if dropped < 1 => { + return Err("command_output_loss records must report dropped_records".to_owned()); + } + Some(COMMAND_OUTPUT_RECORD_TYPE | COMMAND_OUTPUT_LOSS_RECORD_TYPE) => {} + _ => return Err("record_type is unknown".to_owned()), + } + if shape.schema_version != Some(CURRENT_SCHEMA_VERSION) { + return Err(format!("schema_version must be {CURRENT_SCHEMA_VERSION}")); + } + Ok(()) +} diff --git a/crates/sqlbuild-core/src/command_output/mod.rs b/crates/sqlbuild-core/src/command_output/mod.rs new file mode 100644 index 0000000000..88914d5975 --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/mod.rs @@ -0,0 +1,11 @@ +//! Command-output capture: ANSI-free bounded text chunks, buffering and bounded delivery queues. + +pub(crate) mod _helpers; +pub mod constants; +pub mod errors; +pub mod main; +pub mod models; +pub mod types; + +#[cfg(test)] +mod tests; diff --git a/crates/sqlbuild-core/src/command_output/models.rs b/crates/sqlbuild-core/src/command_output/models.rs new file mode 100644 index 0000000000..3edf0fa627 --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/models.rs @@ -0,0 +1,125 @@ +//! Captured text, record drafts and the bounded delivery queue. + +use std::collections::VecDeque; + +use crate::command_output::types::PassText; + +/// The process stream a fragment was written to. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub enum OutputStream { + Stdout, + Stderr, +} + +/// Queue retention class of a record, as `OutputRecordPriority`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RecordPriority { + Bulk, + Terminal, +} + +/// One record the dispatcher turns into a Python `CommandOutputRecord`. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RecordDraft { + pub sequence: u64, + pub stream: OutputStream, + pub message: PassText, + pub chunk_index: usize, + pub chunk_count: usize, + pub priority: RecordPriority, + /// `dropped_records` of a `command_output_loss` record; zero for ordinary output. + pub dropped_records: u64, +} + +/// Line splitting and record buffering state of one dispatcher. +#[derive(Clone, Debug, PartialEq)] +pub struct OutputBuffer { + pub max_record_bytes: i64, + /// Unterminated trailing text per stream: stdout, then stderr. + pub partials: [PassText; 2], + pub buffer_stream: Option, + pub buffer_text: PassText, + pub buffer_bytes: usize, + pub buffer_part_count: usize, + /// Monotonic seconds when the buffered record started. + pub buffer_started_at: Option, + pub sequence: u64, +} + +/// A bounded queue that keeps terminal records over bulk ones, with delivery accounting. +#[derive(Clone, Debug, PartialEq)] +pub struct OutputQueue { + pub items: VecDeque<(T, RecordPriority)>, + pub capacity: usize, + pub accepted: u64, + pub delivered: u64, + pub failed: u64, + pub dropped: u64, +} + +impl OutputBuffer { + /// An empty buffer whose records hold at most `max_record_bytes` UTF-8 bytes. + pub fn new(max_record_bytes: i64) -> Self { + Self { + max_record_bytes, + partials: [PassText::new(), PassText::new()], + buffer_stream: None, + buffer_text: PassText::new(), + buffer_bytes: 0, + buffer_part_count: 0, + buffer_started_at: None, + sequence: 0, + } + } +} + +impl OutputQueue { + /// An empty queue holding at most `capacity` records. + pub fn new(capacity: usize) -> Self { + Self { + items: VecDeque::new(), + capacity, + accepted: 0, + delivered: 0, + failed: 0, + dropped: 0, + } + } +} + +/// The buffer after one call, the records it produced in order, and a failure part-way through. +#[derive(Clone, Debug, PartialEq)] +pub struct BufferOutcome { + pub buffer: OutputBuffer, + pub drafts: Vec, + /// The encode failure Python raised after `drafts` were already enqueued. + pub error: Option, +} + +/// What an idle delivery worker does next. +#[derive(Clone, Debug, PartialEq)] +pub enum IdleStep { + /// Wait for new records, at most this many seconds when a buffered record is pending. + Wait(Option), + /// The buffered record became due and was flushed into these records. + Flushed(Vec), +} + +/// A `CommandOutputRecord`'s fields as the validation sees them; `None` marks a wrong type. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct RecordShape { + /// `invocation_id`, `producer` and `producer_version`, each when it is a `str`. + pub identity: [(&'static str, Option); 3], + /// `run_id`: absent (`None`) or a `str`; `Some(None)` for another type. + pub run_id: Option>, + pub sequence: Option, + pub occurred_at_is_utc: bool, + pub stream_is_known: bool, + pub message_is_text: bool, + pub chunk_index: Option, + pub chunk_count: Option, + pub dropped_records: Option, + pub priority_is_known: bool, + pub record_type: Option, + pub schema_version: Option, +} diff --git a/crates/sqlbuild-core/src/command_output/tests.rs b/crates/sqlbuild-core/src/command_output/tests.rs new file mode 100644 index 0000000000..89648d1562 --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/tests.rs @@ -0,0 +1,6 @@ +#[path = "tests/test_buffering.rs"] +mod buffering; +#[path = "tests/test_types.rs"] +mod test_types; +#[path = "tests/test_text.rs"] +mod text; diff --git a/crates/sqlbuild-core/src/command_output/tests/test_buffering.rs b/crates/sqlbuild-core/src/command_output/tests/test_buffering.rs new file mode 100644 index 0000000000..d509a58d67 --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/tests/test_buffering.rs @@ -0,0 +1,63 @@ +use crate::command_output::main::append_output::append_output; +use crate::command_output::main::close_output::close_output; +use crate::command_output::models::{OutputBuffer, OutputStream}; +use crate::command_output::tests::test_types::BufferingTestCase; + +#[test] +fn given_fragments_when_appending_and_closing_then_lines_coalesce_within_the_record_limit() { + let test_cases = [ + BufferingTestCase { + description: "small lines coalesce into one record at close", + max_record_bytes: 64, + fragments: &["one\ntw", "o\nthree"], + expected_messages: &["one\ntwo\nthree"], + expected_partial: "", + }, + BufferingTestCase { + description: "a line wider than the limit is split into its own chunks", + max_record_bytes: 4, + fragments: &["ab\n", "abcdefgh\n"], + expected_messages: &["ab\n", "abcd", "efgh", "\n"], + expected_partial: "", + }, + ]; + + for test_case in test_cases { + let mut buffer = OutputBuffer::new(test_case.max_record_bytes); + let mut messages: Vec = Vec::new(); + for fragment in test_case.fragments { + let outcome = append_output(buffer, (OutputStream::Stdout, fragment.as_bytes()), 0.0); + buffer = outcome.buffer; + messages.extend( + outcome + .drafts + .iter() + .map(|draft| String::from_utf8_lossy(&draft.message).into_owned()), + ); + } + let closed = close_output(buffer, 1.0); + messages.extend( + closed + .drafts + .iter() + .map(|draft| String::from_utf8_lossy(&draft.message).into_owned()), + ); + + assert_eq!( + ( + messages, + String::from_utf8_lossy(&closed.buffer.partials[0]).into_owned() + ), + ( + test_case + .expected_messages + .iter() + .map(|message| (*message).to_owned()) + .collect::>(), + test_case.expected_partial.to_owned() + ), + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-core/src/command_output/tests/test_text.rs b/crates/sqlbuild-core/src/command_output/tests/test_text.rs new file mode 100644 index 0000000000..251e26cb67 --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/tests/test_text.rs @@ -0,0 +1,64 @@ +use crate::command_output::_helpers::text::{split_lines, strip_ansi}; +use crate::command_output::tests::test_types::TextTestCase; + +#[test] +fn given_terminal_control_sequences_when_stripping_then_only_complete_sequences_are_removed() { + let test_cases = [ + TextTestCase { + description: "colour codes around text", + text: "\u{1b}[1;32mok\u{1b}[0m done", + expected_parts: &["ok done"], + }, + TextTestCase { + description: "hyperlink terminated by BEL and by ST", + text: "\u{1b}]8;;https://example.com\u{7}link\u{1b}]8;;\u{1b}\\", + expected_parts: &["link"], + }, + TextTestCase { + description: "unterminated sequences stay", + text: "\u{1b}]title\u{1b}[12", + expected_parts: &["\u{1b}]title\u{1b}[12"], + }, + ]; + + for test_case in test_cases { + let plain = String::from_utf8(strip_ansi(test_case.text.as_bytes())).unwrap_or_default(); + + assert_eq!( + [plain.as_str()].as_slice(), + test_case.expected_parts, + "{}", + test_case.description + ); + } +} + +#[test] +fn given_mixed_line_boundaries_when_splitting_then_matches_python_splitlines_keepends() { + let test_cases = [ + TextTestCase { + description: "newline, carriage return pair and trailing text", + text: "a\nb\r\nc\rd", + expected_parts: &["a\n", "b\r\n", "c\r", "d"], + }, + TextTestCase { + description: "form feed, file separators and unicode separators", + text: "a\u{c}b\u{1c}c\u{85}d\u{2028}e\u{2029}", + expected_parts: &["a\u{c}", "b\u{1c}", "c\u{85}", "d\u{2028}", "e\u{2029}"], + }, + TextTestCase { + description: "empty text has no lines", + text: "", + expected_parts: &[], + }, + ]; + + for test_case in test_cases { + let lines: Vec = split_lines(test_case.text.as_bytes()) + .into_iter() + .map(|line| String::from_utf8_lossy(line).into_owned()) + .collect(); + + assert_eq!(lines, test_case.expected_parts, "{}", test_case.description); + } +} diff --git a/crates/sqlbuild-core/src/command_output/tests/test_types.rs b/crates/sqlbuild-core/src/command_output/tests/test_types.rs new file mode 100644 index 0000000000..a5d362c295 --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/tests/test_types.rs @@ -0,0 +1,15 @@ +/// One text and the result of a text operation on it. +pub(crate) struct TextTestCase { + pub(crate) description: &'static str, + pub(crate) text: &'static str, + pub(crate) expected_parts: &'static [&'static str], +} + +/// Fragments appended to one dispatcher buffer and the record messages they produce. +pub(crate) struct BufferingTestCase { + pub(crate) description: &'static str, + pub(crate) max_record_bytes: i64, + pub(crate) fragments: &'static [&'static str], + pub(crate) expected_messages: &'static [&'static str], + pub(crate) expected_partial: &'static str, +} diff --git a/crates/sqlbuild-core/src/command_output/types.rs b/crates/sqlbuild-core/src/command_output/types.rs new file mode 100644 index 0000000000..01af513bfc --- /dev/null +++ b/crates/sqlbuild-core/src/command_output/types.rs @@ -0,0 +1,4 @@ +//! Text representations shared by command-output capture. + +/// Text as Python's `str.encode("utf-8", "surrogatepass")` bytes, so lone surrogates survive. +pub type PassText = Vec; diff --git a/crates/sqlbuild-core/src/lib.rs b/crates/sqlbuild-core/src/lib.rs index fce25f3c1c..0a7ad4ba94 100644 --- a/crates/sqlbuild-core/src/lib.rs +++ b/crates/sqlbuild-core/src/lib.rs @@ -9,6 +9,7 @@ pub mod build_progress; pub mod check_command; pub mod cli_text; pub mod clone_text; +pub mod command_output; pub mod compile_timings; pub mod constants; pub mod contract_command; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/cli/entry.rs b/crates/sqlbuild-python/src/bindings/_helpers/cli/entry.rs index f16b619ccb..2f2ffd4833 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/cli/entry.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/cli/entry.rs @@ -17,6 +17,7 @@ use crate::bindings::_helpers::cli::dispatch::{command_facts, dispatch_command}; use crate::bindings::_helpers::cli::launch_state::offer_launch_record; use crate::bindings::_helpers::cli::observability::cli_observability_scope; use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported, within}; +use crate::bindings::_helpers::output_capture::records::invocation_context; use crate::bindings::_helpers::presentation_text::terminal::stdout_supports_color; use sqlbuild_cli::launcher::constants::LAUNCH_RECORD_OPTION; use sqlbuild_core::cli_text::main::format_coded_error::format_coded_error; @@ -30,7 +31,6 @@ const LIFECYCLE_EVENT_MODULE: &str = "sqlbuild.runtime.observability.main.create const INVOCATION_SCOPE_MODULE: &str = "sqlbuild.runtime.observability.main.invocation_scope"; const EXTERNAL_SCOPE_MODULE: &str = "sqlbuild.runtime.observability.main.invocation_external_context_scope"; -const INVOCATION_CONTEXT_MODULE: &str = "sqlbuild.runtime.output_capture.main.invocation_context"; const OUTPUT_CAPTURE_ERRORS_MODULE: &str = "sqlbuild.runtime.output_capture.exceptions"; const DIAGNOSTICS_MODULE: &str = "sqlbuild.diagnostics.main.diagnostics_context"; const TRANSIENT_LINE_MODULE: &str = "sqlbuild._native"; @@ -495,13 +495,7 @@ fn handled_dispatch<'py>(invocation: &Invocation<'py>) -> PyResult) -> PyResult> { - match imported( - py, - INVOCATION_CONTEXT_MODULE, - "invocation_context_from_environment", - )? - .call0() - { + match invocation_context(py, None) { Ok(context) => Ok(context), Err(error) if error.is_instance( diff --git a/crates/sqlbuild-python/src/bindings/_helpers/cli/observability.rs b/crates/sqlbuild-python/src/bindings/_helpers/cli/observability.rs index 6cb27ba298..44a51c59a8 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/cli/observability.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/cli/observability.rs @@ -27,7 +27,7 @@ const EXPORTER_SCOPE_MODULE: &str = "sqlbuild.runtime.event_exporting.main._event_exporter_command_scope"; const EXPORTING_CONSTANTS_MODULE: &str = "sqlbuild.runtime.event_exporting.constants"; const OUTPUT_CAPTURE_MODULE: &str = - "sqlbuild.cli.commands.main.entrypoint._dispatch_with_output_capture"; + "sqlbuild.runtime.output_capture.main.configured_output_capture_scope"; const DEBUG_EVENT_MODULE: &str = "sqlbuild.diagnostics.main.log_debug_event"; const OBSERVABILITY_LOGGER: &str = "sqlbuild.cli.observability"; const DEBUG_COMMAND: &str = "debug"; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/cli/registration.rs b/crates/sqlbuild-python/src/bindings/_helpers/cli/registration.rs index 2857cf1d2b..3a730a4c5b 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/cli/registration.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/cli/registration.rs @@ -10,8 +10,8 @@ use crate::bindings::_helpers::{ audit_command, build_command, build_progress, check, clone, command_runtime, compute_logs, contract_command, cost_command, cost_report, cost_telemetry, dbt_command, dbt_events, dbt_interop, debug_command, diff_report, format_command, freshness, init_command, janitor, - load, plan_command, presentation_text, query_command, result_output, scenario_command, - seed_command, sql_test_command, workspace_commands, + load, output_capture, plan_command, presentation_text, query_command, result_output, + scenario_command, seed_command, sql_test_command, workspace_commands, }; pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { @@ -28,6 +28,7 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { freshness::registration::register(module)?; clone::registration::register(module)?; compute_logs::registration::register(module)?; + output_capture::registration::register(module)?; load::registration::register(module)?; result_output::registration::register(module)?; cost_report::registration::register(module)?; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/mod.rs index 290dead3f9..5d9d284f4b 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/mod.rs @@ -45,6 +45,7 @@ pub(crate) mod load; pub(crate) mod macro_bridge; pub(crate) mod manifest; pub(crate) mod model_config; +pub(crate) mod output_capture; pub(crate) mod plan_command; pub(crate) mod plan_output; pub(crate) mod planner_changes; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/output_capture/context.rs b/crates/sqlbuild-python/src/bindings/_helpers/output_capture/context.rs new file mode 100644 index 0000000000..de6696dc5b --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/output_capture/context.rs @@ -0,0 +1,88 @@ +//! The integration-installed command-output context of the active command. + +use std::sync::Mutex; + +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::{ + Bound, Py, PyAny, PyAnyMethods, PyModule, PyModuleMethods, PyResult, Python, pyclass, + pyfunction, pymethods, wrap_pyfunction, +}; +use pyo3::sync::PyOnceLock; +use pyo3::types::PyTuple; + +use crate::bindings::_helpers::compile_driver::python_calls::{ + call_with, context_variable, imported, +}; +use crate::bindings::_helpers::output_capture::records::MODELS_MODULE; + +const CONTEXT_VARIABLE: &str = "sqlbuild_output_capture_context"; +static CURRENT_CONTEXT: PyOnceLock> = PyOnceLock::new(); + +/// Attaches opaque integration context to command-output records while entered. +#[pyclass(module = "sqlbuild._native", frozen)] +pub(crate) struct CommandOutputContext { + external_context: Py, + token: Mutex>>, +} + +#[pymethods] +impl CommandOutputContext { + fn __enter__(&self, py: Python<'_>) -> PyResult<()> { + let context = call_with( + &imported(py, MODELS_MODULE, "OutputCaptureContext")?, + &[( + "external_context", + imported(py, "builtins", "dict")?.call1((self.external_context.bind(py),))?, + )], + )?; + let token = context_variable(py, &CURRENT_CONTEXT, CONTEXT_VARIABLE)? + .call_method1("set", (context,))?; + *self + .token + .lock() + .map_err(|_| PyRuntimeError::new_err("command output context lock is poisoned"))? = + Some(token.unbind()); + Ok(()) + } + + #[pyo3(signature = (*_exc_info))] + fn __exit__(&self, py: Python<'_>, _exc_info: &Bound<'_, PyTuple>) -> PyResult { + let token = self + .token + .lock() + .map_err(|_| PyRuntimeError::new_err("command output context lock is poisoned"))? + .take(); + if let Some(token) = token { + context_variable(py, &CURRENT_CONTEXT, CONTEXT_VARIABLE)? + .call_method1("reset", (token,))?; + } + Ok(false) + } +} + +/// Attach opaque integration context to command-output records. +#[pyfunction] +#[pyo3(signature = (*, external_context))] +fn command_output_context(external_context: Py) -> CommandOutputContext { + CommandOutputContext { + external_context, + token: Mutex::new(None), + } +} + +/// The optional integration context of the active command. +pub(crate) fn current_context(py: Python<'_>) -> PyResult> { + context_variable(py, &CURRENT_CONTEXT, CONTEXT_VARIABLE)?.call_method0("get") +} + +/// Return optional integration context for the active command. +#[pyfunction] +fn current_output_capture_context(py: Python<'_>) -> PyResult> { + current_context(py) +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_class::()?; + module.add_function(wrap_pyfunction!(command_output_context, module)?)?; + module.add_function(wrap_pyfunction!(current_output_capture_context, module)?) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/output_capture/dispatcher.rs b/crates/sqlbuild-python/src/bindings/_helpers/output_capture/dispatcher.rs new file mode 100644 index 0000000000..d7bd60f0a6 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/output_capture/dispatcher.rs @@ -0,0 +1,512 @@ +//! `OutputCaptureDispatcher`: bounded delivery of captured output under a Python condition. + +use std::collections::HashSet; +use std::sync::{Arc, Mutex, MutexGuard}; +use std::time::Instant; + +use pyo3::exceptions::PyRuntimeError; +use pyo3::prelude::{ + Bound, Py, PyAny, PyAnyMethods, PyErr, PyModule, PyModuleMethods, PyResult, Python, pyclass, + pymethods, +}; +use pyo3::types::{PyDict, PyString, PyTuple}; +use sqlbuild_core::command_output::constants::{ + COMMAND_OUTPUT_LOSS_RECORD_TYPE, MIN_OUTPUT_RECORD_BYTES, MIN_SHUTDOWN_TIMEOUT_SECONDS, +}; +use sqlbuild_core::command_output::main::append_output::append_output; +use sqlbuild_core::command_output::main::close_output::close_output; +use sqlbuild_core::command_output::main::flush_idle_output::flush_idle_output; +use sqlbuild_core::command_output::main::output_loss_summary::output_loss_summary; +use sqlbuild_core::command_output::main::queue_output_record::queue_output_record; +use sqlbuild_core::command_output::models::{ + BufferOutcome, IdleStep, OutputBuffer, OutputQueue, OutputStream, RecordDraft, RecordPriority, +}; + +use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported}; +use crate::bindings::_helpers::output_capture::records::{MODELS_MODULE, TYPES_MODULE}; +use crate::bindings::_helpers::output_capture::values::{ + freeze, object, output_error, pass_text, python_text, unencodable, +}; + +const THREAD_NAME: &str = "sqlbuild-output-exporter"; +const DEFAULT_QUEUE_CAPACITY: i64 = 1024; +const DEFAULT_BATCH_SIZE: i64 = 64; +const DEFAULT_MAX_RECORD_BYTES: i64 = 64 * 1024; +const DEFAULT_SHUTDOWN_TIMEOUT_SECONDS: f64 = 2.0; + +/// Delivery state, read and written only while holding the dispatcher's Python condition. +struct DispatchState { + buffer: OutputBuffer, + queue: OutputQueue>, + accepting: bool, + stopping: bool, + force_stop: bool, + worker_finished: bool, + suppressed_threads: HashSet, +} + +/// What the dispatcher and its worker thread share. +struct Shared { + exporter: Py, + invocation_id: Py, + run_id: Py, + context: Py, + batch_size: usize, + shutdown_timeout_seconds: f64, + clock: Option>, + failure_callback: Option>, + condition: Py, + started: Instant, + state: Mutex, +} + +/// Turn stream fragments into bounded text chunks and export them off execution threads. +#[pyclass(module = "sqlbuild._native", frozen)] +pub(crate) struct OutputCaptureDispatcher { + shared: Arc, + thread: Py, +} + +/// The dispatcher's delivery worker, run as a Python daemon thread's target. +#[pyclass(module = "sqlbuild._native", frozen)] +pub(crate) struct OutputDeliveryWorker { + shared: Arc, +} + +/// The Python condition, entered while alive and exited on drop. +struct Held<'py>(Bound<'py, PyAny>); + +impl Drop for Held<'_> { + fn drop(&mut self) { + let py = self.0.py(); + drop( + self.0 + .call_method1("__exit__", (py.None(), py.None(), py.None())), + ); + } +} + +fn stream_value(stream: OutputStream) -> &'static str { + match stream { + OutputStream::Stdout => "stdout", + OutputStream::Stderr => "stderr", + } +} + +fn priority_value(priority: RecordPriority) -> &'static str { + match priority { + RecordPriority::Bulk => "bulk", + RecordPriority::Terminal => "terminal", + } +} + +fn thread_ident(py: Python<'_>) -> PyResult { + imported(py, "threading", "get_ident")?.call0()?.extract() +} + +impl Shared { + fn locked(&self) -> PyResult> { + self.state + .lock() + .map_err(|_| PyRuntimeError::new_err("output capture state lock is poisoned")) + } + + fn hold<'py>(&self, py: Python<'py>) -> PyResult> { + let condition = self.condition.bind(py).clone(); + condition.call_method0("__enter__")?; + Ok(Held(condition)) + } + + fn now(&self) -> f64 { + self.started.elapsed().as_secs_f64() + } + + fn occurred_at<'py>(&self, py: Python<'py>) -> PyResult> { + match &self.clock { + Some(clock) => clock.bind(py).call0(), + None => imported(py, "datetime", "datetime")? + .call_method1("now", (imported(py, "datetime", "UTC")?,)), + } + } + + /// `CommandOutputRecord` of one draft, timestamped now. + fn record<'py>(&self, py: Python<'py>, draft: &RecordDraft) -> PyResult> { + let mut fields: Vec<(&str, Bound<'py, PyAny>)> = vec![ + ("invocation_id", self.invocation_id.bind(py).clone()), + ("run_id", self.run_id.bind(py).clone()), + ("sequence", object(py, draft.sequence)?), + ("occurred_at", self.occurred_at(py)?), + ( + "stream", + imported(py, TYPES_MODULE, "CommandOutputStream")? + .call1((stream_value(draft.stream),))?, + ), + ("message", python_text(py, &draft.message)?), + ("external_context", self.context.bind(py).clone()), + ("chunk_index", object(py, draft.chunk_index)?), + ("chunk_count", object(py, draft.chunk_count)?), + ( + "priority", + imported(py, TYPES_MODULE, "OutputRecordPriority")? + .call1((priority_value(draft.priority),))?, + ), + ]; + if draft.dropped_records > 0 { + fields.push(("record_type", object(py, COMMAND_OUTPUT_LOSS_RECORD_TYPE)?)); + fields.push(("dropped_records", object(py, draft.dropped_records)?)); + } + call_with( + &imported(py, MODELS_MODULE, "CommandOutputRecord")?, + &fields, + ) + } + + /// Create and queue the records of `drafts` in order. + fn put(&self, py: Python<'_>, drafts: &[RecordDraft]) -> PyResult<()> { + for draft in drafts { + let record = self.record(py, draft)?.unbind(); + let mut state = self.locked()?; + let queue = std::mem::replace(&mut state.queue, OutputQueue::new(0)); + state.queue = queue_output_record(queue, record, draft.priority); + } + Ok(()) + } + + /// Queue an outcome's records, then raise its encode failure. + fn apply(&self, py: Python<'_>, outcome: BufferOutcome, text: &[u8]) -> PyResult<()> { + self.locked()?.buffer = outcome.buffer; + self.put(py, &outcome.drafts)?; + match outcome.error { + Some(error) => Err(unencodable(py, text, &error)), + None => Ok(()), + } + } + + fn take_buffer(&self) -> PyResult { + let mut state = self.locked()?; + let limit = state.buffer.max_record_bytes; + Ok(std::mem::replace( + &mut state.buffer, + OutputBuffer::new(limit), + )) + } + + fn notify(&self, py: Python<'_>, method: &str) -> PyResult<()> { + self.condition.bind(py).call_method0(method)?; + Ok(()) + } + + /// Wait on the condition, at most `timeout` seconds. + fn wait(&self, py: Python<'_>, timeout: Option) -> PyResult<()> { + let condition = self.condition.bind(py); + match timeout { + Some(timeout) => condition.call_method1("wait", (timeout,))?, + None => condition.call_method0("wait")?, + }; + Ok(()) + } + + /// The next batch, or `None` once the worker should stop. + fn next_batch(&self, py: Python<'_>) -> PyResult>>> { + let _held = self.hold(py)?; + loop { + let (idle, stopping) = { + let state = self.locked()?; + (state.queue.items.is_empty(), state.stopping) + }; + if !idle || stopping { + break; + } + let (buffer, step) = flush_idle_output(self.take_buffer()?, self.now()); + self.locked()?.buffer = buffer; + match step { + IdleStep::Wait(timeout) => self.wait(py, timeout)?, + IdleStep::Flushed(drafts) => self.put(py, &drafts)?, + } + } + let mut state = self.locked()?; + if state.force_stop || (state.queue.items.is_empty() && state.stopping) { + return Ok(None); + } + let count = self.batch_size.min(state.queue.items.len()); + Ok(Some( + state + .queue + .items + .drain(..count) + .map(|(record, _)| record) + .collect(), + )) + } + + fn report_failure(&self, py: Python<'_>, error: PyErr) { + if let Some(callback) = &self.failure_callback { + drop(callback.bind(py).call1((error.value(py),))); + } + } + + /// `_export`: deliver one batch with output from this thread excluded from capture. + fn export(&self, py: Python<'_>, batch: Vec>) -> PyResult<()> { + let thread = thread_ident(py)?; + self.locked()?.suppressed_threads.insert(thread); + let size = u64::try_from(batch.len()).unwrap_or(u64::MAX); + let records = PyTuple::new(py, batch)?; + let result = self + .exporter + .bind(py) + .call_method1("export_output", (records,)); + match result { + Ok(_) => { + let _held = self.hold(py)?; + self.locked()?.queue.delivered += size; + } + Err(error) => { + { + let _held = self.hold(py)?; + self.locked()?.queue.failed += size; + } + self.report_failure(py, error); + } + } + self.locked()?.suppressed_threads.remove(&thread); + Ok(()) + } + + fn run(&self, py: Python<'_>) -> PyResult<()> { + let result = (|| -> PyResult<()> { + while let Some(batch) = self.next_batch(py)? { + self.export(py, batch)?; + } + Ok(()) + })(); + let _held = self.hold(py)?; + self.locked()?.worker_finished = true; + self.notify(py, "notify_all")?; + result + } +} + +fn input_error(py: Python<'_>, message: &str) -> PyErr { + output_error(py, "OutputCaptureInputError", message) +} + +#[pymethods] +impl OutputDeliveryWorker { + fn __call__(&self, py: Python<'_>) -> PyResult<()> { + self.shared.run(py) + } +} + +#[pymethods] +impl OutputCaptureDispatcher { + #[new] + #[pyo3(signature = (**options))] + fn new(py: Python<'_>, options: Option>) -> PyResult { + let options = + crate::bindings::_helpers::compile_driver::python_calls::KeywordArguments::read( + py, + "OutputCaptureDispatcher", + options, + ( + &["exporter", "invocation_id"], + &[ + "run_id", + "external_context", + "queue_capacity", + "batch_size", + "max_record_bytes", + "shutdown_timeout_seconds", + "clock", + "failure_callback", + ], + ), + )?; + let number = |name: &str, default: i64| -> PyResult { + let value = options.value(name)?; + if value.is_none() { + Ok(default) + } else { + value.extract() + } + }; + let queue_capacity = number("queue_capacity", DEFAULT_QUEUE_CAPACITY)?; + let batch_size = number("batch_size", DEFAULT_BATCH_SIZE)?; + let max_record_bytes = number("max_record_bytes", DEFAULT_MAX_RECORD_BYTES)?; + let timeout_value = options.value("shutdown_timeout_seconds")?; + let shutdown_timeout_seconds: f64 = if timeout_value.is_none() { + DEFAULT_SHUTDOWN_TIMEOUT_SECONDS + } else { + timeout_value.extract()? + }; + if queue_capacity < 1 || batch_size < 1 { + return Err(input_error( + py, + "output capture queue capacity and batch size must be positive", + )); + } + if max_record_bytes < MIN_OUTPUT_RECORD_BYTES + || shutdown_timeout_seconds < MIN_SHUTDOWN_TIMEOUT_SECONDS + { + return Err(input_error(py, "output capture limits are invalid")); + } + let external = options.value("external_context")?; + let context_source = if external.is_none() || !external.is_truthy()? { + PyDict::new(py).into_any() + } else { + imported(py, "builtins", "dict")?.call1((external,))? + }; + let context = freeze(&context_source, "external_context")?; + let optional = |name: &str| -> PyResult>> { + let value = options.value(name)?; + Ok((!value.is_none()).then(|| value.unbind())) + }; + let shared = Arc::new(Shared { + exporter: options.value("exporter")?.unbind(), + invocation_id: options.value("invocation_id")?.unbind(), + run_id: options.value("run_id")?.unbind(), + context: context.unbind(), + batch_size: usize::try_from(batch_size).unwrap_or(usize::MAX), + shutdown_timeout_seconds, + clock: optional("clock")?, + failure_callback: optional("failure_callback")?, + condition: imported(py, "threading", "Condition")?.call0()?.unbind(), + started: Instant::now(), + state: Mutex::new(DispatchState { + buffer: OutputBuffer::new(max_record_bytes), + queue: OutputQueue::new(usize::try_from(queue_capacity).unwrap_or(usize::MAX)), + accepting: true, + stopping: false, + force_stop: false, + worker_finished: false, + suppressed_threads: HashSet::new(), + }), + }); + let worker = Py::new( + py, + OutputDeliveryWorker { + shared: Arc::clone(&shared), + }, + )?; + let thread = call_with( + &imported(py, "threading", "Thread")?, + &[ + ("target", worker.into_bound(py).into_any()), + ("name", object(py, THREAD_NAME)?), + ("daemon", object(py, true)?), + ], + )?; + thread.call_method0("start")?; + Ok(Self { + shared, + thread: thread.unbind(), + }) + } + + /// Accept one stream fragment without waiting for destination work. + #[pyo3(signature = (*, stream, text))] + fn append( + &self, + py: Python<'_>, + stream: &Bound<'_, PyAny>, + text: &Bound<'_, PyString>, + ) -> PyResult<()> { + let shared = &self.shared; + let thread = thread_ident(py)?; + if text.is_empty()? || shared.locked()?.suppressed_threads.contains(&thread) { + return Ok(()); + } + let stream = match stream.getattr("value")?.extract::()?.as_str() { + "stderr" => OutputStream::Stderr, + _ => OutputStream::Stdout, + }; + let text = pass_text(text)?; + let _held = shared.hold(py)?; + if !shared.locked()?.accepting { + return Ok(()); + } + let outcome = append_output(shared.take_buffer()?, (stream, &text), shared.now()); + let result = shared.apply(py, outcome, &text); + shared.notify(py, "notify")?; + result + } + + /// Flush partial records until a fixed deadline and return without indefinite waits. + fn close<'py>(&self, py: Python<'py>) -> PyResult> { + let shared = &self.shared; + { + let _held = shared.hold(py)?; + if shared.locked()?.accepting { + let outcome = close_output(shared.take_buffer()?, shared.now()); + shared.apply(py, outcome, &[])?; + let summary = { + let mut state = shared.locked()?; + state.accepting = false; + match state.queue.dropped { + 0 => None, + _ => { + let limit = state.buffer.max_record_bytes; + let buffer = + std::mem::replace(&mut state.buffer, OutputBuffer::new(limit)); + let (buffer, draft) = output_loss_summary(buffer, &state.queue); + state.buffer = buffer; + Some(draft) + } + } + }; + if let Some(draft) = summary { + shared.put(py, &[draft])?; + } + shared.locked()?.stopping = true; + shared.notify(py, "notify_all")?; + } + } + let thread = self.thread.bind(py); + call_with( + &thread.getattr("join")?, + &[( + "timeout", + object(py, shared.shutdown_timeout_seconds.max(0.0))?, + )], + )?; + let _held = shared.hold(py)?; + let alive = thread.call_method0("is_alive")?.is_truthy()?; + let counts = { + let mut state = shared.locked()?; + if alive { + state.force_stop = true; + state.queue.dropped += u64::try_from(state.queue.items.len()).unwrap_or(0); + state.queue.items.clear(); + } + ( + state.queue.accepted, + state.queue.delivered, + state.queue.failed, + state.queue.dropped, + state.queue.items.len(), + state.queue.capacity, + state.worker_finished && state.queue.items.is_empty(), + ) + }; + if alive { + shared.notify(py, "notify_all")?; + } + let (accepted, delivered, failed, dropped, depth, capacity, flush_complete) = counts; + call_with( + &imported(py, MODELS_MODULE, "CommandOutputCaptureSummary")?, + &[ + ("accepted", object(py, accepted)?), + ("delivered", object(py, delivered)?), + ("failed", object(py, failed)?), + ("dropped", object(py, dropped)?), + ("queue_depth", object(py, depth)?), + ("queue_capacity", object(py, capacity)?), + ("flush_complete", object(py, flush_complete)?), + ], + ) + } +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_class::()?; + module.add_class::() +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/output_capture/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/output_capture/mod.rs new file mode 100644 index 0000000000..d86f946270 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/output_capture/mod.rs @@ -0,0 +1,6 @@ +pub(crate) mod context; +pub(crate) mod dispatcher; +pub(crate) mod records; +pub(crate) mod registration; +pub(crate) mod text; +pub(crate) mod values; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/output_capture/records.rs b/crates/sqlbuild-python/src/bindings/_helpers/output_capture/records.rs new file mode 100644 index 0000000000..33943740ea --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/output_capture/records.rs @@ -0,0 +1,372 @@ +//! Command-output record validation, wire JSON and the integration's invocation context. + +use pyo3::exceptions::{PyTypeError, PyValueError}; +use pyo3::prelude::{ + Bound, PyAny, PyAnyMethods, PyErr, PyModule, PyModuleMethods, PyResult, Python, pyfunction, + wrap_pyfunction, +}; +use pyo3::types::{PyDict, PyDictMethods, PyString, PyStringMethods}; +use sqlbuild_core::command_output::constants::{ + COMMAND_OUTPUT_LOSS_RECORD_TYPE, COMMAND_OUTPUT_RECORD_TYPE, CURRENT_SCHEMA_VERSION, + INVOCATION_CONTEXT_ENV, MAX_INVOCATION_CONTEXT_BYTES, MAX_INVOCATION_CONTEXT_DEPTH, +}; +use sqlbuild_core::command_output::main::validate_output_record::validate_output_record; +use sqlbuild_core::command_output::models::RecordShape; +use sqlbuild_core::json::main::dumps::dumps; +use sqlbuild_core::json::models::{JsonDialect, JsonValue, StdlibJsonOptions}; + +use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported}; +use crate::bindings::_helpers::output_capture::values::{ + freeze, json_value, object, output_error, plain_int, validation_error, +}; + +pub(crate) const MODELS_MODULE: &str = "sqlbuild.runtime.output_capture.models"; +pub(crate) const TYPES_MODULE: &str = "sqlbuild.runtime.output_capture.types"; +const ENVELOPE_FIELDS: [&str; 16] = [ + "record_id", + "record_type", + "schema_version", + "producer", + "producer_version", + "occurred_at", + "invocation_id", + "run_id", + "sequence", + "stream", + "message", + "chunk_index", + "chunk_count", + "priority", + "dropped_records", + "external_context", +]; +const RECORD_FIELDS: [&str; 15] = [ + "invocation_id", + "run_id", + "sequence", + "occurred_at", + "stream", + "message", + "external_context", + "chunk_index", + "chunk_count", + "priority", + "record_type", + "dropped_records", + "schema_version", + "producer", + "producer_version", +]; + +fn optional_text(value: &Bound<'_, PyAny>) -> PyResult> { + match value.downcast::() { + Ok(text) => Ok(Some(text.to_string_lossy().into_owned())), + Err(_) => Ok(None), + } +} + +fn utc_aware(value: &Bound<'_, PyAny>) -> PyResult { + let py = value.py(); + if !value.is_instance(&imported(py, "datetime", "datetime")?)? + || value.getattr("tzinfo")?.is_none() + { + return Ok(false); + } + let zero = imported(py, "datetime", "timedelta")?.call0()?; + value.call_method0("utcoffset")?.eq(zero) +} + +fn record_shape(record: &Bound<'_, PyAny>) -> PyResult { + let py = record.py(); + let field = |name: &str| record.getattr(name); + let run_id = field("run_id")?; + Ok(RecordShape { + identity: [ + ("invocation_id", optional_text(&field("invocation_id")?)?), + ("producer", optional_text(&field("producer")?)?), + ( + "producer_version", + optional_text(&field("producer_version")?)?, + ), + ], + run_id: if run_id.is_none() { + None + } else { + Some(optional_text(&run_id)?) + }, + sequence: plain_int(&field("sequence")?), + occurred_at_is_utc: utc_aware(&field("occurred_at")?)?, + stream_is_known: field("stream")?.is_instance(&imported( + py, + TYPES_MODULE, + "CommandOutputStream", + )?)?, + message_is_text: field("message")?.is_instance_of::(), + chunk_index: plain_int(&field("chunk_index")?), + chunk_count: plain_int(&field("chunk_count")?), + dropped_records: plain_int(&field("dropped_records")?), + priority_is_known: field("priority")?.is_instance(&imported( + py, + TYPES_MODULE, + "OutputRecordPriority", + )?)?, + record_type: optional_text(&field("record_type")?)?, + schema_version: plain_int(&field("schema_version")?), + }) +} + +/// The frozen `external_context` of a `CommandOutputRecord` that meets the wire contract. +#[pyfunction] +pub(crate) fn freeze_command_output_record<'py>( + record: &Bound<'py, PyAny>, +) -> PyResult> { + let py = record.py(); + validate_output_record(&record_shape(record)?) + .map_err(|message| validation_error(py, &message))?; + let frozen = freeze(&record.getattr("external_context")?, "external_context")?; + if !frozen.is_instance(&imported(py, "collections.abc", "Mapping")?)? { + return Err(validation_error( + py, + "external_context must be a JSON object", + )); + } + Ok(frozen) +} + +/// Serialize one command-output record deterministically. +#[pyfunction] +pub(crate) fn command_output_to_json(record: &Bound<'_, PyAny>) -> PyResult { + let field = |name: &str| record.getattr(name); + let occurred_at: String = call_with( + &field("occurred_at")?.getattr("isoformat")?, + &[("timespec", object(record.py(), "microseconds")?)], + )? + .extract()?; + let mut entries: Vec<(String, JsonValue)> = Vec::new(); + for name in ENVELOPE_FIELDS { + let value: JsonValue = match name { + "occurred_at" => JsonValue::String(occurred_at.replace("+00:00", "Z")), + "stream" | "priority" => json_value(&field(name)?.getattr("value")?)?, + _ => json_value(&field(name)?)?, + }; + entries.push((name.to_owned(), value)); + } + let options = StdlibJsonOptions { + sort_keys: true, + allow_nan: false, + ..StdlibJsonOptions::new(None).with_separators(",", ":") + }; + dumps(&JsonValue::Object(entries), &JsonDialect::Stdlib(options)) + .map_err(|_| PyValueError::new_err("Out of range float values are not JSON compliant")) +} + +fn occurred_at<'py>(value: &Bound<'py, PyAny>) -> PyResult> { + let py = value.py(); + let message = "occurred_at must be an ISO-8601 string"; + let Ok(text) = value.downcast::() else { + return Err(validation_error(py, message)); + }; + let normalized = text.to_cow()?.replace('Z', "+00:00"); + imported(py, "datetime", "datetime")? + .call_method1("fromisoformat", (normalized,)) + .map_err(|error| match error.is_instance_of::(py) { + true => validation_error(py, message), + false => error, + }) +} + +fn record_from_fields<'py>(data: &Bound<'py, PyDict>) -> PyResult> { + let py = data.py(); + let value = |name: &str| -> PyResult> { + data.get_item(name)? + .ok_or_else(|| PyValueError::new_err(name.to_owned())) + }; + let mut fields: Vec<(&str, Bound<'py, PyAny>)> = Vec::new(); + for name in RECORD_FIELDS { + let converted = match name { + "occurred_at" => occurred_at(&value(name)?)?, + "stream" => { + imported(py, TYPES_MODULE, "CommandOutputStream")?.call1((value(name)?,))? + } + "priority" => { + imported(py, TYPES_MODULE, "OutputRecordPriority")?.call1((value(name)?,))? + } + _ => value(name)?, + }; + fields.push((name, converted)); + } + call_with( + &imported(py, MODELS_MODULE, "CommandOutputRecord")?, + &fields, + ) +} + +/// Deserialize one known command-output record strictly. +#[pyfunction] +pub(crate) fn command_output_from_json<'py>( + py: Python<'py>, + raw_json: &Bound<'py, PyAny>, +) -> PyResult> { + let raw = imported(py, "json", "loads")? + .call1((raw_json,)) + .map_err(|error| { + match error.is_instance_of::(py) + || error.is_instance_of::(py) + { + true => validation_error(py, "command output must be valid JSON"), + false => error, + } + })?; + let Ok(data) = raw.downcast::() else { + return Err(validation_error(py, "command output must be a JSON object")); + }; + let mut keys: Vec = Vec::new(); + for key in data.keys() { + keys.push(key.extract()?); + } + keys.sort(); + let mut expected: Vec<&str> = ENVELOPE_FIELDS.to_vec(); + expected.sort_unstable(); + if keys != expected { + return Err(validation_error( + py, + "command output fields do not match schema version 1", + )); + } + let field = |name: &str| -> PyResult> { + data.get_item(name)? + .ok_or_else(|| PyValueError::new_err(name.to_owned())) + }; + if plain_int(&field("schema_version")?) != Some(CURRENT_SCHEMA_VERSION) { + return Err(validation_error( + py, + "unsupported command output schema version", + )); + } + let record_type = optional_text(&field("record_type")?)?; + if !matches!( + record_type.as_deref(), + Some(COMMAND_OUTPUT_RECORD_TYPE | COMMAND_OUTPUT_LOSS_RECORD_TYPE) + ) { + return Err(validation_error(py, "unknown command output record type")); + } + if !field("external_context")?.is_instance_of::() { + return Err(validation_error( + py, + "command output external_context must be an object", + )); + } + let validation_class = imported( + py, + crate::bindings::_helpers::output_capture::values::EXCEPTIONS_MODULE, + "CommandOutputValidationError", + )?; + let record = record_from_fields(data).map_err(|error| { + let typed = + error.is_instance_of::(py) || error.is_instance_of::(py); + if error.is_instance(py, &validation_class) || !typed { + error + } else { + validation_error(py, "command output contains invalid field values") + } + })?; + if !field("record_id")?.eq(record.getattr("record_id")?)? { + return Err(validation_error( + py, + "command output record_id does not match its identity", + )); + } + Ok(record) +} + +fn validate_depth(value: &Bound<'_, PyAny>, depth: usize) -> Result<(), PyErr> { + if depth > MAX_INVOCATION_CONTEXT_DEPTH { + return Err(output_error( + value.py(), + "OutputCaptureInputError", + &format!( + "{INVOCATION_CONTEXT_ENV} must not exceed {MAX_INVOCATION_CONTEXT_DEPTH} levels" + ), + )); + } + if let Ok(mapping) = value.downcast::() { + for item in mapping.values() { + validate_depth(&item, depth + 1)?; + } + } else if value.is_instance_of::() { + for item in value.try_iter()? { + validate_depth(&item?, depth + 1)?; + } + } + Ok(()) +} + +fn input_error(py: Python<'_>, message: &str) -> PyErr { + output_error(py, "OutputCaptureInputError", message) +} + +/// Validated opaque invocation metadata supplied by an integration through the environment. +pub(crate) fn invocation_context<'py>( + py: Python<'py>, + environment: Option>, +) -> PyResult> { + let environment = match environment { + Some(environment) if !environment.is_none() => environment, + _ => imported(py, "os", "environ")?, + }; + let raw = environment.call_method1("get", (INVOCATION_CONTEXT_ENV,))?; + if raw.is_none() { + return Ok(PyDict::new(py).into_any()); + } + let text: String = raw.extract()?; + if text.len() > MAX_INVOCATION_CONTEXT_BYTES { + return Err(input_error( + py, + &format!( + "{INVOCATION_CONTEXT_ENV} must not exceed {MAX_INVOCATION_CONTEXT_BYTES} bytes" + ), + )); + } + let loaded = imported(py, "json", "loads")? + .call1((text,)) + .map_err(|error| match error.is_instance_of::(py) { + true => input_error( + py, + &format!("{INVOCATION_CONTEXT_ENV} must contain valid JSON"), + ), + false => error, + })?; + if !loaded.is_instance_of::() { + return Err(input_error( + py, + &format!("{INVOCATION_CONTEXT_ENV} must contain a JSON object"), + )); + } + validate_depth(&loaded, 1)?; + freeze(&loaded, "external_context").map_err(|error| { + match error.is_instance_of::(py) { + true => input_error(py, &error.value(py).to_string()), + false => error, + } + }) +} + +/// Return validated opaque invocation metadata supplied by an integration. +#[pyfunction] +#[pyo3(signature = (*, environment = None))] +pub(crate) fn invocation_context_from_environment<'py>( + py: Python<'py>, + environment: Option>, +) -> PyResult> { + invocation_context(py, environment) +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(freeze_command_output_record, module)?)?; + module.add_function(wrap_pyfunction!(command_output_to_json, module)?)?; + module.add_function(wrap_pyfunction!(command_output_from_json, module)?)?; + module.add_function(wrap_pyfunction!( + invocation_context_from_environment, + module + )?) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/output_capture/registration.rs b/crates/sqlbuild-python/src/bindings/_helpers/output_capture/registration.rs new file mode 100644 index 0000000000..d23a01a9bd --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/output_capture/registration.rs @@ -0,0 +1,12 @@ +//! Register native command-output capture. + +use pyo3::prelude::{Bound, PyModule, PyResult}; + +use crate::bindings::_helpers::output_capture::{context, dispatcher, records, text}; + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + context::register(module)?; + dispatcher::register(module)?; + records::register(module)?; + text::register(module) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/output_capture/text.rs b/crates/sqlbuild-python/src/bindings/_helpers/output_capture/text.rs new file mode 100644 index 0000000000..e36fb753f7 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/output_capture/text.rs @@ -0,0 +1,31 @@ +//! `chunk_text` for command-output records. + +use pyo3::prelude::{ + Bound, PyAny, PyModule, PyModuleMethods, PyResult, Python, pyfunction, wrap_pyfunction, +}; +use pyo3::types::{PyString, PyTuple}; +use sqlbuild_core::command_output::main::chunk_output_text::chunk_output_text; + +use crate::bindings::_helpers::output_capture::values::{pass_text, python_text, unencodable}; + +/// Split text deterministically without splitting Unicode code points. +#[pyfunction] +#[pyo3(signature = (*, text, max_bytes))] +fn chunk_command_output_text<'py>( + py: Python<'py>, + text: &Bound<'py, PyString>, + max_bytes: i64, +) -> PyResult> { + let bytes = pass_text(text)?; + let chunks = + chunk_output_text(&bytes, max_bytes).map_err(|error| unencodable(py, &bytes, &error))?; + let mut items: Vec> = Vec::new(); + for chunk in &chunks { + items.push(python_text(py, chunk)?); + } + PyTuple::new(py, items) +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(chunk_command_output_text, module)?) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/output_capture/values.rs b/crates/sqlbuild-python/src/bindings/_helpers/output_capture/values.rs new file mode 100644 index 0000000000..d78dee5f8e --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/output_capture/values.rs @@ -0,0 +1,178 @@ +//! Command-output text and JSON values across the Python boundary. + +use pyo3::IntoPyObjectExt; +use pyo3::exceptions::{PyUnicodeEncodeError, PyValueError}; +use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyErr, PyResult, Python}; +use pyo3::types::{ + PyBool, PyBytes, PyBytesMethods, PyDict, PyDictMethods, PyFloat, PyInt, PyList, PyString, + PyStringMethods, PyTuple, PyTypeMethods, +}; +use sqlbuild_core::command_output::errors::UnencodableText; +use sqlbuild_core::command_output::types::PassText; +use sqlbuild_core::json::models::{JsonInteger, JsonValue}; + +use crate::bindings::_helpers::compile_driver::python_calls::imported; +use crate::bindings::_helpers::macro_bridge::store_keys::text; + +pub(crate) const EXCEPTIONS_MODULE: &str = "sqlbuild.runtime.output_capture.exceptions"; +const SURROGATE_PASS: &str = "surrogatepass"; +const UTF8: &str = "utf-8"; + +/// The exception class `name` of the output capture exceptions module, raised with `message`. +pub(crate) fn output_error(py: Python<'_>, name: &str, message: &str) -> PyErr { + match imported(py, EXCEPTIONS_MODULE, name).and_then(|class| class.call1((message,))) { + Ok(error) => PyErr::from_value(error), + Err(error) => error, + } +} + +/// `CommandOutputValidationError(message)`. +pub(crate) fn validation_error(py: Python<'_>, message: &str) -> PyErr { + output_error(py, "CommandOutputValidationError", message) +} + +/// A Python `str` as `surrogatepass` UTF-8 bytes. +pub(crate) fn pass_text(value: &Bound<'_, PyString>) -> PyResult { + if let Ok(text) = value.to_str() { + return Ok(text.as_bytes().to_vec()); + } + Ok(value + .call_method1("encode", (UTF8, SURROGATE_PASS))? + .downcast_into::()? + .as_bytes() + .to_vec()) +} + +/// `surrogatepass` UTF-8 bytes as a Python `str`. +pub(crate) fn python_text<'py>(py: Python<'py>, value: &[u8]) -> PyResult> { + match std::str::from_utf8(value) { + Ok(text) => Ok(PyString::new(py, text).into_any()), + Err(_) => PyBytes::new(py, value).call_method1("decode", (UTF8, SURROGATE_PASS)), + } +} + +/// The `UnicodeEncodeError` Python raises encoding `text` with a lone surrogate. +pub(crate) fn unencodable(py: Python<'_>, text: &[u8], error: &UnencodableText) -> PyErr { + let raised = python_text(py, text).and_then(|text| { + py.get_type::().call1(( + UTF8, + text, + error.position, + error.position + 1, + "surrogates not allowed", + )) + }); + match raised { + Ok(exception) => PyErr::from_value(exception), + Err(failure) => failure, + } +} + +fn type_name(value: &Bound<'_, PyAny>) -> PyResult { + Ok(value.get_type().name()?.to_string()) +} + +/// `freeze_command_output_json`: validate and recursively freeze a JSON-compatible value. +pub(crate) fn freeze<'py>(value: &Bound<'py, PyAny>, path: &str) -> PyResult> { + let py: Python<'py> = value.py(); + if value.is_none() + || value.is_instance_of::() + || value.is_instance_of::() + || value.is_instance_of::() + { + return Ok(value.clone()); + } + if value.is_instance_of::() { + if !value.extract::()?.is_finite() { + return Err(validation_error( + py, + &format!("{path} must not contain NaN or infinity"), + )); + } + return Ok(value.clone()); + } + if value.is_instance(&imported(py, "collections.abc", "Mapping")?)? { + let frozen: Bound<'py, PyDict> = PyDict::new(py); + for item in value.call_method0("items")?.try_iter()? { + let (key, item): (Bound<'py, PyAny>, Bound<'py, PyAny>) = item?.extract()?; + let Ok(key_text) = key.downcast::() else { + return Err(validation_error( + py, + &format!("{path} keys must be strings"), + )); + }; + let child = format!("{path}.{}", key_text.to_cow()?); + frozen.set_item(&key, freeze(&item, &child)?)?; + } + return imported(py, "types", "MappingProxyType")?.call1((frozen,)); + } + if value.is_instance_of::() || value.is_instance_of::() { + let mut items: Vec> = Vec::new(); + for (index, item) in value.try_iter()?.enumerate() { + items.push(freeze(&item?, &format!("{path}[{index}]"))?); + } + return Ok(PyTuple::new(py, items)?.into_any()); + } + Err(validation_error( + py, + &format!( + "{path} contains non-JSON value of type {}", + type_name(value)? + ), + )) +} + +/// A frozen command-output JSON value for the native emitter. +pub(crate) fn json_value(value: &Bound<'_, PyAny>) -> PyResult { + if value.is_none() { + return Ok(JsonValue::Null); + } + if let Ok(flag) = value.downcast::() { + return Ok(JsonValue::Bool(flag.is_truthy()?)); + } + if value.is_instance_of::() { + let decimal: String = value.str()?.extract()?; + return JsonInteger::parse(&decimal) + .map(JsonValue::Integer) + .ok_or_else(|| PyValueError::new_err("command output integer is not JSON")); + } + if value.is_instance_of::() { + return Ok(JsonValue::Float(value.extract::()?)); + } + if let Ok(string) = value.downcast::() { + return text(string); + } + if value.is_instance_of::() || value.is_instance_of::() { + let mut items: Vec = Vec::new(); + for item in value.try_iter()? { + items.push(json_value(&item?)?); + } + return Ok(JsonValue::Array(items)); + } + let mut entries: Vec<(String, JsonValue)> = Vec::new(); + for item in value.call_method0("items")?.try_iter()? { + let (key, item): (String, Bound<'_, PyAny>) = item?.extract()?; + entries.push((key, json_value(&item)?)); + } + Ok(JsonValue::Object(entries)) +} + +/// `value` as an `int` other than `bool`, saturated to `i64`, or `None`. +pub(crate) fn plain_int(value: &Bound<'_, PyAny>) -> Option { + if value.is_instance_of::() || !value.is_instance_of::() { + return None; + } + let saturated = match value.gt(0) { + Ok(true) => i64::MAX, + _ => i64::MIN, + }; + Some(value.extract::().unwrap_or(saturated)) +} + +/// A Rust value as a Python object. +pub(crate) fn object<'py, T: IntoPyObjectExt<'py>>( + py: Python<'py>, + value: T, +) -> PyResult> { + value.into_bound_py_any(py) +} diff --git a/python-boundary.toml b/python-boundary.toml index a80f8da9de..cd96621e9b 100644 --- a/python-boundary.toml +++ b/python-boundary.toml @@ -280,9 +280,39 @@ reason = "executor" note = "Exports execution events while the executor runs." [[stays]] -path = "runtime/output_capture" -reason = "executor" -note = "Captures stdout and stderr of running Python nodes for the executor." +path = "runtime/output_capture/classes" +reason = "user-code" +note = "Rule 1 capture shim: sys.stdout/sys.stderr tees forwarding accepted writes from user code, adapters and the executor to the native OutputCaptureDispatcher." + +[[stays]] +path = "runtime/output_capture/main" +reason = "user-code" +note = "Rule 1 capture shim entry: swaps sys.stdout/sys.stderr for the tees around one CLI command (Rust cli/observability.rs)." + +[[stays]] +path = "runtime/output_capture/__init__.py" +reason = "public-api" +note = "Package marker of the public command-output shapes below." + +[[stays]] +path = "runtime/output_capture/models.py" +reason = "public-api" +note = "Rule 3 shapes: CommandOutputRecord and sink definitions re-exported by sqlbuild.sinks; validation is native validate_command_output_record." + +[[stays]] +path = "runtime/output_capture/types.py" +reason = "public-api" +note = "Rule 3 shapes: CommandOutputStream re-exported by sqlbuild.sinks, priorities and the exporter protocol." + +[[stays]] +path = "runtime/output_capture/exceptions.py" +reason = "public-api" +note = "Rule 3 shapes: CommandOutputValidationError re-exported by sqlbuild.sinks and raised by the native records." + +[[stays]] +path = "runtime/output_capture/constants.py" +reason = "public-api" +note = "Rule 3 shapes: the invocation-context variable and wire constants used by the models and integrations." [[stays]] path = "integrations/dagster" @@ -319,11 +349,6 @@ path = "cli/commands/main/entrypoint/_dispatch_with_compute_logs.py" reason = "user-code" note = "Captures the command's stdout, stderr and logging into the local compute log around user code." -[[stays]] -path = "cli/commands/main/entrypoint/_dispatch_with_output_capture.py" -reason = "user-code" -note = "Tees the command's Python output streams to configured output capture sinks." - [[stays]] path = "cli/commands/_helpers/entry/compute_log_diagnostics.py" reason = "user-code" diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index 46672425f5..362c97f69e 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -355,6 +355,12 @@ from sqlbuild.runtime.contracts.models import ConnectionHooks from sqlbuild.runtime.contracts.types import ExecutionResourceKind from sqlbuild.runtime.observability.classes.event_dispatcher import EventDispatcher from sqlbuild.runtime.observability.models import DiagnosticLog, LifecycleEvent +from sqlbuild.runtime.output_capture.models import ( + CommandOutputCaptureSummary, + CommandOutputRecord, + OutputCaptureContext, +) +from sqlbuild.runtime.output_capture.types import CommandOutputStream from sqlbuild.spec.contracts.models import ( CostConfig, DbtConfig, @@ -3452,6 +3458,43 @@ class CliObservabilityScope: def cli_observability_scope(*, args: CliNamespace, project_dir: Path) -> CliObservabilityScope: ... +class OutputCaptureDispatcher: + """Turn stream fragments into bounded text chunks and export them off execution threads.""" + + def __init__( + self, + *, + exporter: object, + invocation_id: str, + run_id: str | None = None, + external_context: Mapping[str, object] | None = None, + queue_capacity: int = 1024, + batch_size: int = 64, + max_record_bytes: int = 65536, + shutdown_timeout_seconds: float = 2.0, + clock: Callable[[], datetime] | None = None, + failure_callback: Callable[[BaseException], object] | None = None, + ) -> None: ... + def append(self, *, stream: CommandOutputStream, text: str) -> None: ... + def close(self) -> CommandOutputCaptureSummary: ... + +class OutputDeliveryWorker: + def __call__(self) -> None: ... + +class CommandOutputContext: + def __enter__(self) -> None: ... + def __exit__(self, *exc_info: object) -> bool: ... + +def command_output_context(*, external_context: Mapping[str, object]) -> CommandOutputContext: ... +def current_output_capture_context() -> OutputCaptureContext | None: ... +def chunk_command_output_text(*, text: str, max_bytes: int) -> tuple[str, ...]: ... +def freeze_command_output_record(record: CommandOutputRecord, /) -> Mapping[str, object]: ... +def command_output_to_json(record: CommandOutputRecord, /) -> str: ... +def command_output_from_json(raw_json: str, /) -> CommandOutputRecord: ... +def invocation_context_from_environment( + *, environment: Mapping[str, str] | None = None +) -> Mapping[str, object]: ... + class LocalFilesystemComputeLogStorage: """Ordered fsync capture storage subject to underlying filesystem and hardware guarantees.""" diff --git a/src/sqlbuild/runtime/output_capture/classes/binary_tee.py b/src/sqlbuild/runtime/output_capture/classes/binary_tee.py index d81e5911e3..36bca44213 100644 --- a/src/sqlbuild/runtime/output_capture/classes/binary_tee.py +++ b/src/sqlbuild/runtime/output_capture/classes/binary_tee.py @@ -4,7 +4,7 @@ from typing import Any -from sqlbuild.runtime.output_capture.classes.dispatcher import OutputCaptureDispatcher +from sqlbuild._native import OutputCaptureDispatcher from sqlbuild.runtime.output_capture.types import CommandOutputStream diff --git a/src/sqlbuild/runtime/output_capture/classes/text_tee.py b/src/sqlbuild/runtime/output_capture/classes/text_tee.py index 821086df56..6fd94e2f6b 100644 --- a/src/sqlbuild/runtime/output_capture/classes/text_tee.py +++ b/src/sqlbuild/runtime/output_capture/classes/text_tee.py @@ -4,8 +4,8 @@ from typing import Any +from sqlbuild._native import OutputCaptureDispatcher from sqlbuild.runtime.output_capture.classes.binary_tee import BinaryOutputTee -from sqlbuild.runtime.output_capture.classes.dispatcher import OutputCaptureDispatcher from sqlbuild.runtime.output_capture.types import CommandOutputStream diff --git a/src/sqlbuild/runtime/output_capture/main/configured_output_capture_scope.py b/src/sqlbuild/runtime/output_capture/main/configured_output_capture_scope.py index 014f8d7db2..f20ef69ab9 100644 --- a/src/sqlbuild/runtime/output_capture/main/configured_output_capture_scope.py +++ b/src/sqlbuild/runtime/output_capture/main/configured_output_capture_scope.py @@ -1,4 +1,4 @@ -"""Configured command-output capture for one CLI command.""" +"""Tee the process streams of one CLI command into configured command-output sinks.""" from __future__ import annotations @@ -7,15 +7,13 @@ from contextlib import contextmanager from typing import Any -from sqlbuild.runtime.observability.models import ExecutionIdentity -from sqlbuild.runtime.output_capture.classes.dispatcher import OutputCaptureDispatcher -from sqlbuild.runtime.output_capture.classes.text_tee import TextOutputTee -from sqlbuild.runtime.output_capture.main.current_output_capture_context import ( +from sqlbuild._native import ( + OutputCaptureDispatcher, current_output_capture_context, -) -from sqlbuild.runtime.output_capture.main.invocation_context import ( invocation_context_from_environment, ) +from sqlbuild.runtime.observability.models import ExecutionIdentity +from sqlbuild.runtime.output_capture.classes.text_tee import TextOutputTee from sqlbuild.runtime.output_capture.models import OutputCaptureContext from sqlbuild.runtime.output_capture.types import ( CommandOutputBatchExporter, diff --git a/src/sqlbuild/runtime/output_capture/models.py b/src/sqlbuild/runtime/output_capture/models.py index 320d2382d7..7675bcda3b 100644 --- a/src/sqlbuild/runtime/output_capture/models.py +++ b/src/sqlbuild/runtime/output_capture/models.py @@ -4,16 +4,14 @@ from collections.abc import Callable, Mapping from dataclasses import dataclass, field -from datetime import UTC, datetime +from datetime import datetime -from sqlbuild.runtime.output_capture._helpers.json import freeze_command_output_json +from sqlbuild._native import freeze_command_output_record from sqlbuild.runtime.output_capture.constants import ( - COMMAND_OUTPUT_LOSS_RECORD_TYPE, COMMAND_OUTPUT_RECORD_TYPE, CURRENT_COMMAND_OUTPUT_SCHEMA_VERSION, SQLBUILD_COMMAND_OUTPUT_PRODUCER_VERSION, ) -from sqlbuild.runtime.output_capture.exceptions import CommandOutputValidationError from sqlbuild.runtime.output_capture.types import CommandOutputStream, OutputRecordPriority @@ -40,70 +38,7 @@ class CommandOutputRecord: def __post_init__(self) -> None: """Validate and freeze the known command-output wire contract.""" - for field_name in ("invocation_id", "producer", "producer_version"): - value: object = getattr(self, field_name) - if not isinstance(value, str) or not value: - raise CommandOutputValidationError(f"{field_name} must be a non-empty string") - if self.run_id is not None and (not isinstance(self.run_id, str) or not self.run_id): - raise CommandOutputValidationError("run_id must be null or a non-empty string") - if ( - isinstance(self.sequence, bool) - or not isinstance(self.sequence, int) - or self.sequence < 0 - ): - raise CommandOutputValidationError("sequence must be a non-negative integer") - if ( - not isinstance(self.occurred_at, datetime) - or self.occurred_at.tzinfo is None - or self.occurred_at.utcoffset() != UTC.utcoffset(self.occurred_at) - ): - raise CommandOutputValidationError("occurred_at must be a timezone-aware UTC datetime") - if not isinstance(self.stream, CommandOutputStream): - raise CommandOutputValidationError("stream must be stdout or stderr") - if not isinstance(self.message, str): - raise CommandOutputValidationError("message must be a string") - for field_name, minimum in ( - ("chunk_index", 0), - ("chunk_count", 1), - ("dropped_records", 0), - ): - value = getattr(self, field_name) - if isinstance(value, bool) or not isinstance(value, int) or value < minimum: - raise CommandOutputValidationError( - f"{field_name} must be an integer of at least {minimum}" - ) - if self.chunk_index >= self.chunk_count: - raise CommandOutputValidationError("chunk_index must be less than chunk_count") - if not isinstance(self.priority, OutputRecordPriority): - raise CommandOutputValidationError("priority must be bulk or terminal") - if not isinstance(self.record_type, str) or self.record_type not in { - COMMAND_OUTPUT_RECORD_TYPE, - COMMAND_OUTPUT_LOSS_RECORD_TYPE, - }: - raise CommandOutputValidationError("record_type is unknown") - if self.record_type == COMMAND_OUTPUT_RECORD_TYPE and self.dropped_records != 0: - raise CommandOutputValidationError( - "command_output records must have zero dropped_records" - ) - if self.record_type == COMMAND_OUTPUT_LOSS_RECORD_TYPE and self.dropped_records < 1: - raise CommandOutputValidationError( - "command_output_loss records must report dropped_records" - ) - if ( - isinstance(self.schema_version, bool) - or not isinstance(self.schema_version, int) - or self.schema_version != CURRENT_COMMAND_OUTPUT_SCHEMA_VERSION - ): - raise CommandOutputValidationError( - f"schema_version must be {CURRENT_COMMAND_OUTPUT_SCHEMA_VERSION}" - ) - frozen_context: object = freeze_command_output_json( - value=self.external_context, - path="external_context", - ) - if not isinstance(frozen_context, Mapping): - raise CommandOutputValidationError("external_context must be a JSON object") - object.__setattr__(self, "external_context", frozen_context) + object.__setattr__(self, "external_context", freeze_command_output_record(self)) @property def record_id(self) -> str: diff --git a/src/sqlbuild/sinks.py b/src/sqlbuild/sinks.py index a96160f1d0..c0445e039a 100644 --- a/src/sqlbuild/sinks.py +++ b/src/sqlbuild/sinks.py @@ -1,5 +1,8 @@ """Public typed sink declarations for lifecycle events and command output.""" +from sqlbuild._native import command_output_context as command_output_context +from sqlbuild._native import command_output_from_json as command_output_from_json +from sqlbuild._native import command_output_to_json as command_output_to_json from sqlbuild.runtime.event_exporting.main.command_output_sink import ( command_output_sink as command_output_sink, ) @@ -23,15 +26,6 @@ from sqlbuild.runtime.output_capture.exceptions import ( CommandOutputValidationError as CommandOutputValidationError, ) -from sqlbuild.runtime.output_capture.main.command_output_context import ( - command_output_context as command_output_context, -) -from sqlbuild.runtime.output_capture.main.command_output_from_json import ( - command_output_from_json as command_output_from_json, -) -from sqlbuild.runtime.output_capture.main.command_output_to_json import ( - command_output_to_json as command_output_to_json, -) from sqlbuild.runtime.output_capture.models import ( CommandOutputRecord as CommandOutputRecord, ) diff --git a/tests/unit/src/sqlbuild/cli/commands/main/entrypoint/test_dispatch_with_output_capture.py b/tests/unit/src/sqlbuild/cli/commands/main/entrypoint/test_dispatch_with_output_capture.py index b374f64a28..6e8fe08857 100644 --- a/tests/unit/src/sqlbuild/cli/commands/main/entrypoint/test_dispatch_with_output_capture.py +++ b/tests/unit/src/sqlbuild/cli/commands/main/entrypoint/test_dispatch_with_output_capture.py @@ -7,12 +7,12 @@ import pytest -from sqlbuild.cli.commands.main.entrypoint._dispatch_with_output_capture import ( - configured_output_capture_scope, -) from sqlbuild.observability import ExecutionIdentity from sqlbuild.runtime.event_exporting.classes.command_scope import EventExporterCommandScope from sqlbuild.runtime.output_capture.constants import INVOCATION_CONTEXT_ENV +from sqlbuild.runtime.output_capture.main.configured_output_capture_scope import ( + configured_output_capture_scope, +) from sqlbuild.sinks import ( CommandOutputRecord, CommandOutputValidationError, diff --git a/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/test_text.py b/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/test_text.py index b32863bf60..3970d41580 100644 --- a/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/test_text.py +++ b/tests/unit/src/sqlbuild/runtime/output_capture/_helpers/test_text.py @@ -4,7 +4,7 @@ import pytest -from sqlbuild.runtime.output_capture._helpers.text import chunk_text +from sqlbuild._native import chunk_command_output_text as chunk_text from tests.unit.src.sqlbuild.runtime.output_capture._helpers._test_types import ( ChunkTextErrorTestCase, ChunkTextFunction, diff --git a/tests/unit/src/sqlbuild/runtime/output_capture/classes/helpers.py b/tests/unit/src/sqlbuild/runtime/output_capture/classes/helpers.py index 382409cb3e..22856d5dcb 100644 --- a/tests/unit/src/sqlbuild/runtime/output_capture/classes/helpers.py +++ b/tests/unit/src/sqlbuild/runtime/output_capture/classes/helpers.py @@ -6,7 +6,7 @@ from collections.abc import Callable, Mapping from datetime import UTC, datetime -from sqlbuild.runtime.output_capture.classes.dispatcher import OutputCaptureDispatcher +from sqlbuild._native import OutputCaptureDispatcher from sqlbuild.runtime.output_capture.models import CommandOutputRecord diff --git a/tests/unit/src/sqlbuild/runtime/output_capture/classes/test_output_capture_dispatcher.py b/tests/unit/src/sqlbuild/runtime/output_capture/classes/test_output_capture_dispatcher.py index e4a907bdb5..c371e52fe7 100644 --- a/tests/unit/src/sqlbuild/runtime/output_capture/classes/test_output_capture_dispatcher.py +++ b/tests/unit/src/sqlbuild/runtime/output_capture/classes/test_output_capture_dispatcher.py @@ -9,7 +9,7 @@ import pytest -from sqlbuild.runtime.output_capture.classes.dispatcher import OutputCaptureDispatcher +from sqlbuild._native import OutputCaptureDispatcher from sqlbuild.runtime.output_capture.classes.text_tee import TextOutputTee from sqlbuild.runtime.output_capture.models import CommandOutputCaptureSummary from sqlbuild.runtime.output_capture.types import CommandOutputStream diff --git a/tests/unit/src/sqlbuild/runtime/output_capture/main/test_invocation_context.py b/tests/unit/src/sqlbuild/runtime/output_capture/main/test_invocation_context.py index 7d5928e28e..19a5031ccd 100644 --- a/tests/unit/src/sqlbuild/runtime/output_capture/main/test_invocation_context.py +++ b/tests/unit/src/sqlbuild/runtime/output_capture/main/test_invocation_context.py @@ -7,11 +7,9 @@ import pytest +from sqlbuild._native import invocation_context_from_environment from sqlbuild.runtime.output_capture.constants import INVOCATION_CONTEXT_ENV from sqlbuild.runtime.output_capture.exceptions import OutputCaptureInputError -from sqlbuild.runtime.output_capture.main.invocation_context import ( - invocation_context_from_environment, -) from tests.unit.src.sqlbuild.runtime.output_capture.main._test_types import ( InvocationContextTestCase, ) From 1de6e4c1098c45e190ec58c78cd3a774ffb48f5f Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 15 Oct 2026 11:55:48 +0100 Subject: [PATCH 11/12] refactor: host execution history backends natively --- crates/README.md | 2 +- .../src/bindings/_helpers/cli/registration.rs | 7 +- .../_helpers/execution_history/appends.rs | 36 + .../_helpers/execution_history/conversion.rs | 456 ++++++++ .../_helpers/execution_history/functions.rs | 134 +++ .../_helpers/execution_history/mod.rs | 7 + .../_helpers/execution_history/postgres.rs | 987 ++++++++++++++++++ .../execution_history/registration.rs | 13 + .../_helpers/execution_history/sqlite.rs | 921 ++++++++++++++++ .../_helpers/execution_history/statements.rs | 70 ++ .../src/bindings/_helpers/mod.rs | 1 + .../execution_history/_helpers/base64url.rs | 70 ++ .../execution_history/_helpers/dialects.rs | 68 ++ .../src/execution_history/_helpers/filters.rs | 40 + .../src/execution_history/_helpers/mod.rs | 5 + .../src/execution_history/_helpers/schema.rs | 16 + .../execution_history/_helpers/timestamps.rs | 29 + .../src/execution_history/constants.rs | 136 +++ .../src/execution_history/errors.rs | 31 + .../main/canonical_event_id.rs | 15 + .../main/decoded_run_cursor.rs | 39 + .../execution_history/main/event_cursor.rs | 9 + .../main/event_cursor_position.rs | 36 + .../main/event_insert_query.rs | 45 + .../main/event_page_query.rs | 61 ++ .../src/execution_history/main/mod.rs | 20 + .../execution_history/main/prepared_events.rs | 39 + .../execution_history/main/project_runs.rs | 111 ++ .../src/execution_history/main/run_cursor.rs | 31 + .../main/run_insert_query.rs | 37 + .../execution_history/main/run_page_query.rs | 74 ++ .../main/schema_migration_needed.rs | 25 + .../main/sqlite_database_path.rs | 42 + .../execution_history/main/validate_filter.rs | 54 + .../execution_history/main/validate_page.rs | 25 + .../main/validate_page_limit.rs | 16 + .../main/validate_postgres_settings.rs | 62 ++ .../main/validate_schema_target.rs | 14 + .../main/validate_stored_content.rs | 25 + .../main/validate_stored_event.rs | 28 + .../src/execution_history/mod.rs | 11 + .../src/execution_history/models.rs | 178 ++++ .../src/execution_history/tests.rs | 12 + .../src/execution_history/tests/helpers.rs | 65 ++ .../execution_history/tests/test_cursors.rs | 128 +++ .../tests/test_projection.rs | 77 ++ .../execution_history/tests/test_queries.rs | 68 ++ .../src/execution_history/tests/test_types.rs | 55 + .../tests/test_validation.rs | 143 +++ .../src/execution_history/types.rs | 6 + crates/sqlbuild-state/src/lib.rs | 1 + fensu.toml | 5 + python-boundary.toml | 24 +- scripts/fensu_policy/constants.py | 3 - .../rules/execution_observability.py | 3 +- src/sqlbuild/_native.pyi | 107 ++ src/sqlbuild/execution_history.py | 22 +- src/sqlbuild/postgres_history.py | 6 +- .../execution_history/_helpers/__init__.py | 1 - .../_helpers/canonical_events.py | 24 - .../execution_history/_helpers/projection.py | 94 -- .../execution_history/_helpers/validation.py | 38 - .../execution_history/classes/__init__.py | 1 - .../classes/postgres_execution_history.py | 868 --------------- .../classes/sqlite_execution_history.py | 795 -------------- .../execution_history/main/__init__.py | 1 - .../main/append_and_project.py | 23 - .../main/canonical_event_content.py | 12 - .../main/canonical_event_id.py | 12 - .../execution_history/main/project_runs.py | 14 - .../main/validate_page_limit.py | 11 - .../runtime/execution_history/models.py | 67 +- src/sqlbuild/sqlite_history.py | 6 +- .../test_sqlite_execution_history.py | 18 +- 74 files changed, 4745 insertions(+), 1991 deletions(-) create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/execution_history/appends.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/execution_history/conversion.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/execution_history/functions.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/execution_history/mod.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/execution_history/postgres.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/execution_history/registration.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/execution_history/sqlite.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/execution_history/statements.rs create mode 100644 crates/sqlbuild-state/src/execution_history/_helpers/base64url.rs create mode 100644 crates/sqlbuild-state/src/execution_history/_helpers/dialects.rs create mode 100644 crates/sqlbuild-state/src/execution_history/_helpers/filters.rs create mode 100644 crates/sqlbuild-state/src/execution_history/_helpers/mod.rs create mode 100644 crates/sqlbuild-state/src/execution_history/_helpers/schema.rs create mode 100644 crates/sqlbuild-state/src/execution_history/_helpers/timestamps.rs create mode 100644 crates/sqlbuild-state/src/execution_history/constants.rs create mode 100644 crates/sqlbuild-state/src/execution_history/errors.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/canonical_event_id.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/decoded_run_cursor.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/event_cursor.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/event_cursor_position.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/event_insert_query.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/event_page_query.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/mod.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/prepared_events.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/project_runs.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/run_cursor.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/run_insert_query.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/run_page_query.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/schema_migration_needed.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/sqlite_database_path.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/validate_filter.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/validate_page.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/validate_page_limit.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/validate_postgres_settings.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/validate_schema_target.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/validate_stored_content.rs create mode 100644 crates/sqlbuild-state/src/execution_history/main/validate_stored_event.rs create mode 100644 crates/sqlbuild-state/src/execution_history/mod.rs create mode 100644 crates/sqlbuild-state/src/execution_history/models.rs create mode 100644 crates/sqlbuild-state/src/execution_history/tests.rs create mode 100644 crates/sqlbuild-state/src/execution_history/tests/helpers.rs create mode 100644 crates/sqlbuild-state/src/execution_history/tests/test_cursors.rs create mode 100644 crates/sqlbuild-state/src/execution_history/tests/test_projection.rs create mode 100644 crates/sqlbuild-state/src/execution_history/tests/test_queries.rs create mode 100644 crates/sqlbuild-state/src/execution_history/tests/test_types.rs create mode 100644 crates/sqlbuild-state/src/execution_history/tests/test_validation.rs create mode 100644 crates/sqlbuild-state/src/execution_history/types.rs delete mode 100644 src/sqlbuild/runtime/execution_history/_helpers/__init__.py delete mode 100644 src/sqlbuild/runtime/execution_history/_helpers/canonical_events.py delete mode 100644 src/sqlbuild/runtime/execution_history/_helpers/projection.py delete mode 100644 src/sqlbuild/runtime/execution_history/_helpers/validation.py delete mode 100644 src/sqlbuild/runtime/execution_history/classes/__init__.py delete mode 100644 src/sqlbuild/runtime/execution_history/classes/postgres_execution_history.py delete mode 100644 src/sqlbuild/runtime/execution_history/classes/sqlite_execution_history.py delete mode 100644 src/sqlbuild/runtime/execution_history/main/__init__.py delete mode 100644 src/sqlbuild/runtime/execution_history/main/append_and_project.py delete mode 100644 src/sqlbuild/runtime/execution_history/main/canonical_event_content.py delete mode 100644 src/sqlbuild/runtime/execution_history/main/canonical_event_id.py delete mode 100644 src/sqlbuild/runtime/execution_history/main/project_runs.py delete mode 100644 src/sqlbuild/runtime/execution_history/main/validate_page_limit.py diff --git a/crates/README.md b/crates/README.md index 54922e3765..52618c87d0 100644 --- a/crates/README.md +++ b/crates/README.md @@ -8,7 +8,7 @@ that can be tested and benchmarked with `cargo test` alone. |---|---| | `sqlbuild-core` | Constants shared by several crates, the panic boundary that turns an unwinding native failure into a named compiler error, passive identity, diagnostic and SQL value types, text positions, a JSON emitter whose output matches Python's `json.dumps` and `orjson` byte for byte, and run-cost allocation with exact Python `decimal` default-context arithmetic and Snowflake query-history SQL. | | `sqlbuild-cursor` | Microbatch cursor algebra: typed cursor values, grain flooring and stepping, half-open interval clamping, merging and batch splitting, timestamp text normalization and calendar durations, with Python `datetime` results and errors. Zones other than fixed offsets are resolved through a caller-supplied lookup. | -| `sqlbuild-state` | Append-only warehouse state rows: content-addressed migration event IDs, typed event SQL literals and statements, and source freshness data-version normalization, hashing, equivalence, age policy and latest-state read SQL, byte-compatible with the rows Python wrote. | +| `sqlbuild-state` | Append-only warehouse state rows: content-addressed migration event IDs, typed event SQL literals and statements, and source freshness data-version normalization, hashing, equivalence, age policy and latest-state read SQL, byte-compatible with the rows Python wrote. It also owns execution history: lifecycle event batch digests and conflicts, run projection, event and run cursors, filter and page validation, and the SQLite and PostgreSQL statements the native history backends run. | | `sqlbuild-pysource` | Static reading of project Python source with the ruff parser: macro module structure, literal SQL arguments, definition locations and string-literal edits. It never imports or runs user code; CPython stays the authority on whether a file is valid Python. | | `sqlbuild-cache` | The shared native store every native cache uses (one atomically replaced, checksummed file per cache kind, keyed by content digests) and the one fingerprinting scheme for cache keys: the content digest and the project file fingerprint. | | `sqlbuild-sqltext` | Lexical SQL text without polyglot: comment, quote and parenthesis scanning, the rules quote policy, model header tokenization and matching, static variable substitution, static reference extraction, and the general reference scan with table-function call arguments under each adapter's lexical rules. | diff --git a/crates/sqlbuild-python/src/bindings/_helpers/cli/registration.rs b/crates/sqlbuild-python/src/bindings/_helpers/cli/registration.rs index 3a730a4c5b..c09f1aea65 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/cli/registration.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/cli/registration.rs @@ -9,9 +9,9 @@ use crate::bindings::_helpers::cli::{ use crate::bindings::_helpers::{ audit_command, build_command, build_progress, check, clone, command_runtime, compute_logs, contract_command, cost_command, cost_report, cost_telemetry, dbt_command, dbt_events, - dbt_interop, debug_command, diff_report, format_command, freshness, init_command, janitor, - load, output_capture, plan_command, presentation_text, query_command, result_output, - scenario_command, seed_command, sql_test_command, workspace_commands, + dbt_interop, debug_command, diff_report, execution_history, format_command, freshness, + init_command, janitor, load, output_capture, plan_command, presentation_text, query_command, + result_output, scenario_command, seed_command, sql_test_command, workspace_commands, }; pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { @@ -28,6 +28,7 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { freshness::registration::register(module)?; clone::registration::register(module)?; compute_logs::registration::register(module)?; + execution_history::registration::register(module)?; output_capture::registration::register(module)?; load::registration::register(module)?; result_output::registration::register(module)?; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/execution_history/appends.rs b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/appends.rs new file mode 100644 index 0000000000..1238a733fe --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/appends.rs @@ -0,0 +1,36 @@ +//! Canonical identity, content and envelopes of events about to be appended. + +use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyResult}; +use sqlbuild_state::execution_history::main::canonical_event_id::canonical_event_id; +use sqlbuild_state::execution_history::main::prepared_events::prepared_events; +use sqlbuild_state::execution_history::models::{EventSubmission, PreparedEvent}; + +use crate::bindings::_helpers::execution_history::conversion::{ + checked, envelope_of, event_id_of, event_json, +}; + +/// The required stable id of a known or opaque event. +pub(crate) fn event_id(event: &Bound<'_, PyAny>) -> PyResult { + let raw: Option = event_id_of(event)?; + checked(event.py(), canonical_event_id(raw.as_deref())) +} + +/// Deterministic canonical content of an event with a stable id. +pub(crate) fn event_content(event: &Bound<'_, PyAny>) -> PyResult { + event_id(event)?; + event_json(event) +} + +/// Every event in `events` digested, rejecting ids reused with different content in the batch. +pub(crate) fn prepared_batch(events: &Bound<'_, PyAny>) -> PyResult> { + let mut submissions: Vec = Vec::new(); + for event in events.try_iter()? { + let event: Bound<'_, PyAny> = event?; + submissions.push(EventSubmission { + event_id: event_id(&event)?, + content: event_content(&event)?, + envelope: envelope_of(&event)?, + }); + } + checked(events.py(), prepared_events(submissions)) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/execution_history/conversion.rs b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/conversion.rs new file mode 100644 index 0000000000..5e51a40232 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/conversion.rs @@ -0,0 +1,456 @@ +//! Execution history models, timestamps, SQL parameters and errors across the Python boundary. + +use pyo3::IntoPyObjectExt; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyErr, PyResult, Python}; +use pyo3::types::{PyBool, PyInt, PyList, PyString, PyStringMethods, PyTuple}; +use sqlbuild_state::execution_history::constants::ENVELOPE_TEXT_FIELDS; +use sqlbuild_state::execution_history::errors::{HistoryError, HistoryErrorKind}; +use sqlbuild_state::execution_history::models::{ + EventEnvelope, EventFilterSpec, FilterTime, ProjectionFact, RunFact, RunFilterSpec, RunRow, + RunStatus, SqlParam, +}; +use sqlbuild_state::execution_history::types::UtcTime; + +use crate::bindings::_helpers::compile_driver::python_calls::{call_with, imported}; +use crate::bindings::_helpers::cursor::values::{DATETIME_MODULE, local_date_time}; + +pub(crate) const EXCEPTIONS_MODULE: &str = "sqlbuild.runtime.execution_history.exceptions"; +pub(crate) const MODELS_MODULE: &str = "sqlbuild.runtime.execution_history.models"; +const TYPES_MODULE: &str = "sqlbuild.runtime.execution_history.types"; +const EVENT_MODELS_MODULE: &str = "sqlbuild.runtime.observability.models"; +const EVENT_TO_JSON_MODULE: &str = "sqlbuild.runtime.observability.main.lifecycle_event_to_json"; +const EVENT_FROM_JSON_MODULE: &str = + "sqlbuild.runtime.observability.main.lifecycle_event_from_json"; +const UTC_SUFFIX: &str = "Z"; +const UTC_OFFSET: &str = "+00:00"; +const DEFAULT_SCHEMA_VERSION: i64 = 1; +const UTC_OFFSET_SECONDS: f64 = 0.0; + +fn exception_name(kind: HistoryErrorKind) -> &'static str { + match kind { + HistoryErrorKind::Storage => "ExecutionHistoryStorageError", + HistoryErrorKind::IntegrityConflict => "IntegrityConflictError", + HistoryErrorKind::InvalidEvent => "InvalidEventError", + HistoryErrorKind::InvalidCursor => "InvalidCursorError", + HistoryErrorKind::InvalidFilter => "InvalidFilterError", + HistoryErrorKind::InvalidLimit => "InvalidLimitError", + HistoryErrorKind::UnsupportedSchemaVersion => "UnsupportedSchemaVersionError", + HistoryErrorKind::ProjectionConsistency => "ProjectionConsistencyError", + } +} + +/// The execution history exception for a rejection of `kind`. +pub(crate) fn rejection(py: Python<'_>, kind: HistoryErrorKind, message: &str) -> PyErr { + match imported(py, EXCEPTIONS_MODULE, exception_name(kind)) + .and_then(|class| class.call1((message,))) + { + Ok(error) => PyErr::from_value(error), + Err(error) => error, + } +} + +/// An `ExecutionHistoryStorageError`. +pub(crate) fn storage_error(py: Python<'_>, message: &str) -> PyErr { + rejection(py, HistoryErrorKind::Storage, message) +} + +/// An `ExecutionHistoryStorageError` raised `from cause`. +pub(crate) fn storage_error_from(py: Python<'_>, message: &str, cause: PyErr) -> PyErr { + let error: PyErr = storage_error(py, message); + error.set_cause(py, Some(cause)); + error +} + +/// `result`, with a history rejection raised as its Python exception. +pub(crate) fn checked(py: Python<'_>, result: Result) -> PyResult { + result.map_err(|error| rejection(py, error.kind, &error.message)) +} + +/// Whether `error` is an instance of the execution history exception named `name`. +pub(crate) fn history_error_of_kind(py: Python<'_>, error: &PyErr, name: &str) -> PyResult { + Ok(error.is_instance(py, &imported(py, EXCEPTIONS_MODULE, name)?)) +} + +/// `value` when it is an `int` other than `bool` that fits, as `isinstance(value, int)` checks. +pub(crate) fn exact_int(value: &Bound<'_, PyAny>) -> Option { + if !value.is_instance_of::() || value.is_instance_of::() { + return None; + } + value.extract().unwrap_or(None) +} + +/// `value` as text when it is a `str`. +pub(crate) fn text_of(value: &Bound<'_, PyAny>) -> PyResult> { + match value.downcast::() { + Ok(text) => Ok(Some(text.to_str()?.to_owned())), + Err(_) => Ok(None), + } +} + +/// `value` as text, or `None` when it is `None`. +pub(crate) fn optional_text(value: &Bound<'_, PyAny>) -> PyResult> { + if value.is_none() { + return Ok(None); + } + Ok(Some(value.extract()?)) +} + +/// `datetime.UTC`. +pub(crate) fn utc(py: Python<'_>) -> PyResult> { + imported(py, DATETIME_MODULE, "UTC") +} + +/// An aware `datetime` as its UTC wall time. +pub(crate) fn utc_time(value: &Bound<'_, PyAny>) -> PyResult { + local_date_time(&value.call_method1("astimezone", (utc(value.py())?,))?) +} + +/// A UTC `datetime` at `value`. +pub(crate) fn python_time<'py>(py: Python<'py>, value: &UtcTime) -> PyResult> { + imported(py, DATETIME_MODULE, "datetime")?.call1(( + value.date.year, + value.date.month, + value.date.day, + value.hour, + value.minute, + value.second, + value.microsecond, + utc(py)?, + )) +} + +/// `datetime.now(UTC)`. +pub(crate) fn now(py: Python<'_>) -> PyResult { + let current = imported(py, DATETIME_MODULE, "datetime")?.call_method1("now", (utc(py)?,))?; + local_date_time(¤t) +} + +/// Whether a `datetime`'s `tzinfo` is exactly `UTC`. +pub(crate) fn utc_zone(value: &Bound<'_, PyAny>) -> PyResult { + Ok(value.getattr("tzinfo")?.is(&utc(value.py())?)) +} + +/// `datetime.fromisoformat(text.replace("Z", "+00:00"))`, or `None` on `ValueError`. +pub(crate) fn parsed_iso<'py>(py: Python<'py>, text: &str) -> PyResult>> { + let normalized: String = text.replace(UTC_SUFFIX, UTC_OFFSET); + match imported(py, DATETIME_MODULE, "datetime")?.call_method1("fromisoformat", (normalized,)) { + Ok(value) => Ok(Some(value)), + Err(error) if error.is_instance_of::(py) => Ok(None), + Err(error) => Err(error), + } +} + +fn filter_time(value: &Bound<'_, PyAny>) -> PyResult> { + if value.is_none() { + return Ok(None); + } + Ok(Some(FilterTime { + instant: utc_time(value)?, + is_utc: utc_zone(value)?, + })) +} + +fn texts(value: &Bound<'_, PyAny>) -> PyResult> { + value + .try_iter()? + .map(|item| item.and_then(|item| item.extract())) + .collect() +} + +/// An `EventFilter`'s fields. +pub(crate) fn event_filter_spec(value: &Bound<'_, PyAny>) -> PyResult { + let family: Bound<'_, PyAny> = value.getattr("family")?; + Ok(EventFilterSpec { + invocation_id: optional_text(&value.getattr("invocation_id")?)?, + run_id: optional_text(&value.getattr("run_id")?)?, + producer: optional_text(&value.getattr("producer")?)?, + event_types: texts(&value.getattr("event_types")?)?, + family: if family.is_none() { + None + } else { + Some(family.getattr("value")?.extract()?) + }, + occurred_at_start: filter_time(&value.getattr("occurred_at_start")?)?, + occurred_at_end: filter_time(&value.getattr("occurred_at_end")?)?, + }) +} + +/// A `RunStatus` member's status. +fn run_status(value: &Bound<'_, PyAny>) -> PyResult { + let text: String = value.getattr("value")?.extract()?; + parsed_status(value.py(), &text) +} + +/// The status named by stored `text`, as `RunStatus(text)` accepts it. +pub(crate) fn parsed_status(py: Python<'_>, text: &str) -> PyResult { + let known: Option = [RunStatus::Unknown, RunStatus::Completed, RunStatus::Failed] + .into_iter() + .find(|status| status.value() == text); + if let Some(status) = known { + return Ok(status); + } + imported(py, TYPES_MODULE, "RunStatus")?.call1((text,))?; + Ok(RunStatus::Unknown) +} + +/// A `RunFilter`'s fields. +pub(crate) fn run_filter_spec(value: &Bound<'_, PyAny>) -> PyResult { + let statuses: Vec = value + .getattr("statuses")? + .try_iter()? + .map(|item| item.and_then(|item| run_status(&item))) + .collect::>()?; + Ok(RunFilterSpec { + invocation_id: optional_text(&value.getattr("invocation_id")?)?, + statuses, + created_at_start: filter_time(&value.getattr("created_at_start")?)?, + created_at_end: filter_time(&value.getattr("created_at_end")?)?, + }) +} + +/// Whether `event` is a known-schema `LifecycleEvent`. +pub(crate) fn known_event(event: &Bound<'_, PyAny>) -> PyResult { + event.is_instance(&imported( + event.py(), + EVENT_MODELS_MODULE, + "LifecycleEvent", + )?) +} + +/// The `event_id` a known event carries or an opaque envelope holds as a string. +pub(crate) fn event_id_of(event: &Bound<'_, PyAny>) -> PyResult> { + if known_event(event)? { + return text_of(&event.getattr("event_id")?); + } + text_of(&raw_field(event, "event_id")?) +} + +fn raw_field<'py>(event: &Bound<'py, PyAny>, name: &str) -> PyResult> { + event.getattr("raw")?.call_method1("get", (name,)) +} + +/// `lifecycle_event_to_json(event)`. +pub(crate) fn event_json(event: &Bound<'_, PyAny>) -> PyResult { + imported(event.py(), EVENT_TO_JSON_MODULE, "lifecycle_event_to_json")? + .call1((event,))? + .extract() +} + +/// `lifecycle_event_from_json(payload)`. +pub(crate) fn event_from_json<'py>(payload: &Bound<'py, PyAny>) -> PyResult> { + imported( + payload.py(), + EVENT_FROM_JSON_MODULE, + "lifecycle_event_from_json", + )? + .call1((payload,)) +} + +/// An opaque envelope's `occurred_at` when it is ISO text with a zero UTC offset. +fn opaque_occurred_at(event: &Bound<'_, PyAny>) -> PyResult> { + let py: Python<'_> = event.py(); + let Some(text) = text_of(&raw_field(event, "occurred_at")?)? else { + return Ok(None); + }; + let Some(parsed) = parsed_iso(py, &text)? else { + return Ok(None); + }; + let offset: Bound<'_, PyAny> = parsed.call_method0("utcoffset")?; + if offset.is_none() + || offset.call_method0("total_seconds")?.extract::()? != UTC_OFFSET_SECONDS + { + return Ok(None); + } + Ok(Some(utc_time(&parsed)?)) +} + +/// The stable envelope fields stored next to an event. +pub(crate) fn envelope_of(event: &Bound<'_, PyAny>) -> PyResult { + let mut envelope = EventEnvelope::default(); + if known_event(event)? { + envelope.schema_version = event.getattr("schema_version")?.extract()?; + for (slot, name) in envelope.text_fields.iter_mut().zip(ENVELOPE_TEXT_FIELDS) { + *slot = optional_text(&event.getattr(name)?)?; + } + envelope.occurred_at = Some(utc_time(&event.getattr("occurred_at")?)?); + return Ok(envelope); + } + envelope.schema_version = + exact_int(&raw_field(event, "schema_version")?).unwrap_or(DEFAULT_SCHEMA_VERSION); + for (slot, name) in envelope.text_fields.iter_mut().zip(ENVELOPE_TEXT_FIELDS) { + *slot = text_of(&raw_field(event, name)?)?.filter(|text| !text.is_empty()); + } + envelope.occurred_at = opaque_occurred_at(event)?; + Ok(envelope) +} + +/// A `StoredEvent` as projection reads it. +pub(crate) fn projection_fact(stored: &Bound<'_, PyAny>) -> PyResult { + let event: Bound<'_, PyAny> = stored.getattr("event")?; + let mut run: Option = None; + if known_event(&event)? { + let run_id: Bound<'_, PyAny> = event.getattr("run_id")?; + if !run_id.is_none() { + run = Some(RunFact { + run_id: run_id.extract()?, + invocation_id: event.getattr("invocation_id")?.extract()?, + event_type: event.getattr("event_type")?.extract()?, + occurred_at: utc_time(&event.getattr("occurred_at")?)?, + }); + } + } + Ok(ProjectionFact { + storage_order: stored.getattr("storage_order")?.extract()?, + cursor: stored.getattr("cursor")?.extract()?, + received_at: utc_time(&stored.getattr("received_at")?)?, + run, + }) +} + +/// A `RunRecord` as a projected run. +pub(crate) fn run_row(record: &Bound<'_, PyAny>) -> PyResult { + let optional_time = |name: &str| -> PyResult> { + let value: Bound<'_, PyAny> = record.getattr(name)?; + if value.is_none() { + return Ok(None); + } + utc_time(&value).map(Some) + }; + Ok(RunRow { + run_id: record.getattr("run_id")?.extract()?, + invocation_id: record.getattr("invocation_id")?.extract()?, + created_at: utc_time(&record.getattr("created_at")?)?, + status: run_status(&record.getattr("status")?)?, + is_complete: record.getattr("is_complete")?.is_truthy()?, + last_event_cursor: record.getattr("last_event_cursor")?.extract()?, + last_storage_order: record.getattr("last_storage_order")?.extract()?, + command: optional_text(&record.getattr("command")?)?, + target: optional_text(&record.getattr("target")?)?, + environment: optional_text(&record.getattr("environment")?)?, + started_at: optional_time("started_at")?, + ended_at: optional_time("ended_at")?, + }) +} + +fn optional_python_time<'py>( + py: Python<'py>, + value: Option<&UtcTime>, +) -> PyResult> { + match value { + Some(value) => python_time(py, value), + None => Ok(py.None().into_bound(py)), + } +} + +/// The `RunRecord` of a projected run. +pub(crate) fn run_record<'py>(py: Python<'py>, run: &RunRow) -> PyResult> { + let status = imported(py, TYPES_MODULE, "RunStatus")?.call1((run.status.value(),))?; + call_with( + &imported(py, MODELS_MODULE, "RunRecord")?, + &[ + ("run_id", run.run_id.clone().into_bound_py_any(py)?), + ( + "invocation_id", + run.invocation_id.clone().into_bound_py_any(py)?, + ), + ("created_at", python_time(py, &run.created_at)?), + ("status", status), + ("is_complete", run.is_complete.into_bound_py_any(py)?), + ( + "last_event_cursor", + run.last_event_cursor.clone().into_bound_py_any(py)?, + ), + ( + "last_storage_order", + run.last_storage_order.into_bound_py_any(py)?, + ), + ("command", run.command.clone().into_bound_py_any(py)?), + ("target", run.target.clone().into_bound_py_any(py)?), + ( + "environment", + run.environment.clone().into_bound_py_any(py)?, + ), + ( + "started_at", + optional_python_time(py, run.started_at.as_ref())?, + ), + ("ended_at", optional_python_time(py, run.ended_at.as_ref())?), + ], + ) +} + +/// A tuple of the `RunRecord`s of projected runs. +pub(crate) fn run_records<'py>(py: Python<'py>, runs: &[RunRow]) -> PyResult> { + let records: Vec> = runs + .iter() + .map(|run| run_record(py, run)) + .collect::>()?; + PyTuple::new(py, records) +} + +/// The projected runs of an iterable of `RunRecord`s. +pub(crate) fn run_rows(records: &Bound<'_, PyAny>) -> PyResult> { + records + .try_iter()? + .map(|record| record.and_then(|record| run_row(&record))) + .collect() +} + +/// The projection facts of an iterable of `StoredEvent`s. +pub(crate) fn projection_facts(stored: &Bound<'_, PyAny>) -> PyResult> { + stored + .try_iter()? + .map(|item| item.and_then(|item| projection_fact(&item))) + .collect() +} + +/// A `StoredEvent` decoded from its durable columns. +pub(crate) fn stored_event<'py>( + py: Python<'py>, + (storage_order, cursor): (i64, String), + received_at: Bound<'py, PyAny>, + payload: &Bound<'py, PyAny>, +) -> PyResult> { + call_with( + &imported(py, MODELS_MODULE, "StoredEvent")?, + &[ + ("storage_order", storage_order.into_bound_py_any(py)?), + ("cursor", cursor.into_bound_py_any(py)?), + ("received_at", received_at), + ("event", event_from_json(payload)?), + ], + ) +} + +/// A SQL parameter as the driver binds it. +pub(crate) fn python_param<'py>(py: Python<'py>, param: &SqlParam) -> PyResult> { + match param { + SqlParam::Null => Ok(py.None().into_bound(py)), + SqlParam::Text(text) => text.clone().into_bound_py_any(py), + SqlParam::Integer(value) => value.into_bound_py_any(py), + SqlParam::Flag(value) => value.into_bound_py_any(py), + SqlParam::Timestamp(value) => python_time(py, value), + SqlParam::TextList(values) => Ok(PyList::new(py, values)?.into_any()), + } +} + +/// The parameters of a query as a tuple. +pub(crate) fn python_params<'py>( + py: Python<'py>, + params: &[SqlParam], +) -> PyResult> { + let values: Vec> = params + .iter() + .map(|param| python_param(py, param)) + .collect::>()?; + PyTuple::new(py, values) +} + +/// `uuid.uuid4().hex`. +pub(crate) fn new_namespace(py: Python<'_>) -> PyResult { + imported(py, "uuid", "uuid4")? + .call0()? + .getattr("hex")? + .extract() +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/execution_history/functions.rs b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/functions.rs new file mode 100644 index 0000000000..8f3bb313bd --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/functions.rs @@ -0,0 +1,134 @@ +//! Backend-neutral execution history entry points and model validation. + +use pyo3::prelude::{ + Bound, PyAny, PyAnyMethods, PyErr, PyModule, PyModuleMethods, PyResult, Python, pyfunction, + wrap_pyfunction, +}; +use pyo3::types::PyTuple; +use sqlbuild_state::execution_history::main::project_runs::project_runs; +use sqlbuild_state::execution_history::main::validate_filter::validate_filter; +use sqlbuild_state::execution_history::main::validate_page::validate_page; +use sqlbuild_state::execution_history::main::validate_page_limit::validate_page_limit; +use sqlbuild_state::execution_history::main::validate_stored_event::validate_stored_event; +use sqlbuild_state::execution_history::models::FilterSpec; + +use crate::bindings::_helpers::execution_history::appends::{event_content, event_id}; +use crate::bindings::_helpers::execution_history::conversion::{ + checked, event_filter_spec, exact_int, projection_facts, run_filter_spec, run_records, + run_rows, text_of, utc_zone, +}; + +/// Return the required stable ID from a known or opaque lifecycle event. +#[pyfunction] +fn canonical_event_id(event: &Bound<'_, PyAny>) -> PyResult { + event_id(event) +} + +/// Return deterministic content for lifecycle event idempotency comparison. +#[pyfunction] +fn canonical_event_content(event: &Bound<'_, PyAny>) -> PyResult { + event_content(event) +} + +/// Apply durable event facts in storage order to immutable run projections. +#[pyfunction] +#[pyo3(signature = (*, stored_events, current_runs = None))] +fn project_execution_runs<'py>( + py: Python<'py>, + stored_events: &Bound<'py, PyAny>, + current_runs: Option<&Bound<'py, PyAny>>, +) -> PyResult> { + let facts = projection_facts(stored_events)?; + let current = match current_runs { + Some(current_runs) => run_rows(current_runs)?, + None => Vec::new(), + }; + run_records(py, &checked(py, project_runs(facts, current))?) +} + +/// Validate a positive page limit no greater than the public maximum. +#[pyfunction] +fn validate_history_page_limit(limit: &Bound<'_, PyAny>) -> Result<(), PyErr> { + checked(limit.py(), validate_page_limit(exact_int(limit))) +} + +/// Append facts durably before publishing their rebuildable run projection. +#[pyfunction] +#[pyo3(signature = (*, event_log, run_storage, events))] +fn append_and_project<'py>( + event_log: &Bound<'py, PyAny>, + run_storage: &Bound<'py, PyAny>, + events: &Bound<'py, PyAny>, +) -> PyResult> { + let stored: Bound<'py, PyAny> = event_log.call_method1("append_events", (events,))?; + run_storage.call_method1("project", (&stored,))?; + Ok(stored) +} + +/// Validate a `StoredEvent`'s position, cursor and UTC receipt time. +#[pyfunction] +fn validate_stored_history_event(record: &Bound<'_, PyAny>) -> Result<(), PyErr> { + let py: Python<'_> = record.py(); + let storage_order: i64 = record.getattr("storage_order")?.extract()?; + let cursor: String = text_of(&record.getattr("cursor")?)?.unwrap_or_default(); + let received_at_is_utc: bool = utc_zone(&record.getattr("received_at")?)?; + checked( + py, + validate_stored_event(storage_order, &cursor, received_at_is_utc), + ) +} + +/// Validate an `EventFilter`. +#[pyfunction] +fn validate_history_event_filter(event_filter: &Bound<'_, PyAny>) -> Result<(), PyErr> { + checked( + event_filter.py(), + validate_filter(FilterSpec::Event(&event_filter_spec(event_filter)?)), + ) +} + +/// Validate a `RunFilter`. +#[pyfunction] +fn validate_history_run_filter(run_filter: &Bound<'_, PyAny>) -> Result<(), PyErr> { + checked( + run_filter.py(), + validate_filter(FilterSpec::Run(&run_filter_spec(run_filter)?)), + ) +} + +/// Validate an `EventPage` or `RunPage` next cursor against its records. +#[pyfunction] +#[pyo3(signature = (*, page, is_event_page))] +fn validate_history_page(page: &Bound<'_, PyAny>, is_event_page: bool) -> Result<(), PyErr> { + let py: Python<'_> = page.py(); + let records: Bound<'_, PyAny> = page.getattr("records")?; + let has_records: bool = records.is_truthy()?; + let next_cursor: Bound<'_, PyAny> = page.getattr("next_cursor")?; + let next_text: Option = text_of(&next_cursor)?; + let last_text: Option = if is_event_page && has_records { + text_of(&records.get_item(-1)?.getattr("cursor")?)? + } else { + None + }; + let next: Option<&str> = if next_cursor.is_none() { + None + } else { + Some(next_text.as_deref().unwrap_or_default()) + }; + checked( + py, + validate_page(is_event_page, has_records, (next, last_text.as_deref())), + ) +} + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_function(wrap_pyfunction!(canonical_event_id, module)?)?; + module.add_function(wrap_pyfunction!(canonical_event_content, module)?)?; + module.add_function(wrap_pyfunction!(project_execution_runs, module)?)?; + module.add_function(wrap_pyfunction!(validate_history_page_limit, module)?)?; + module.add_function(wrap_pyfunction!(append_and_project, module)?)?; + module.add_function(wrap_pyfunction!(validate_stored_history_event, module)?)?; + module.add_function(wrap_pyfunction!(validate_history_event_filter, module)?)?; + module.add_function(wrap_pyfunction!(validate_history_run_filter, module)?)?; + module.add_function(wrap_pyfunction!(validate_history_page, module)?) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/execution_history/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/mod.rs new file mode 100644 index 0000000000..10306a6285 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/mod.rs @@ -0,0 +1,7 @@ +pub(crate) mod appends; +pub(crate) mod conversion; +pub(crate) mod functions; +pub(crate) mod postgres; +pub(crate) mod registration; +pub(crate) mod sqlite; +pub(crate) mod statements; diff --git a/crates/sqlbuild-python/src/bindings/_helpers/execution_history/postgres.rs b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/postgres.rs new file mode 100644 index 0000000000..b55b674307 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/postgres.rs @@ -0,0 +1,987 @@ +//! `PostgresExecutionHistory`: deployed PostgreSQL event-log and run-projection storage. + +use std::sync::OnceLock; +use std::sync::atomic::{AtomicBool, AtomicI64, Ordering}; + +use pyo3::IntoPyObjectExt; +use pyo3::exceptions::PyException; +use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyErr, PyResult, Python, pyclass, pymethods}; +use pyo3::types::{PyBool, PyDateTime, PyDict, PyTuple}; +use sqlbuild_state::execution_history::constants::{ + DEFAULT_PAGE_LIMIT, POSTGRES_ADVISORY_LOCK, POSTGRES_ALL_EVENTS, POSTGRES_ALL_RUNS, + POSTGRES_CONNECT_TIMEOUT_SECONDS, POSTGRES_DELETE_RUNS, POSTGRES_EVENT_BY_ID, + POSTGRES_EVENT_EXISTS, POSTGRES_HISTORY_LOCK_ID, POSTGRES_MIGRATION_LOCK_ID, + POSTGRES_MIGRATIONS_EXIST, POSTGRES_NAMESPACE, POSTGRES_PING, POSTGRES_RECORD_MIGRATION, + POSTGRES_RETRYABLE_SQLSTATES, POSTGRES_RUN_BY_ID, POSTGRES_RUN_EXISTS, POSTGRES_SCHEMA, + POSTGRES_SCHEMA_VERSION, POSTGRES_SESSION_TIMEOUT, POSTGRES_STATEMENT_TIMEOUT_MS, + POSTGRES_TRANSACTION_RETRIES, POSTGRES_TRANSACTION_TIMEOUT, RUN_CURSOR_TIMESTAMP, + SCHEMA_VERSION, UNKNOWN_EVENT_CURSOR, UNKNOWN_RUN_CURSOR, +}; +use sqlbuild_state::execution_history::errors::HistoryErrorKind; +use sqlbuild_state::execution_history::main::decoded_run_cursor::decoded_run_cursor; +use sqlbuild_state::execution_history::main::event_cursor::event_cursor; +use sqlbuild_state::execution_history::main::event_cursor_position::event_cursor_position; +use sqlbuild_state::execution_history::main::event_insert_query::event_insert_query; +use sqlbuild_state::execution_history::main::event_page_query::event_page_query; +use sqlbuild_state::execution_history::main::project_runs::project_runs; +use sqlbuild_state::execution_history::main::run_cursor::run_cursor; +use sqlbuild_state::execution_history::main::run_insert_query::run_insert_query; +use sqlbuild_state::execution_history::main::run_page_query::run_page_query; +use sqlbuild_state::execution_history::main::schema_migration_needed::schema_migration_needed; +use sqlbuild_state::execution_history::main::validate_page_limit::validate_page_limit; +use sqlbuild_state::execution_history::main::validate_postgres_settings::validate_postgres_settings; +use sqlbuild_state::execution_history::main::validate_schema_target::validate_schema_target; +use sqlbuild_state::execution_history::main::validate_stored_content::validate_stored_content; +use sqlbuild_state::execution_history::models::{ + Dialect, PostgresSettings, PreparedEvent, Query, RunCursorKey, RunRow, SqlParam, +}; +use sqlbuild_state::execution_history::types::UtcTime; + +use crate::bindings::_helpers::compile_driver::python_calls::{ + KeywordArguments, call_with, imported, keywords, within, +}; +use crate::bindings::_helpers::execution_history::appends::prepared_batch; +use crate::bindings::_helpers::execution_history::conversion::{ + MODELS_MODULE, checked, event_filter_spec, exact_int, history_error_of_kind, new_namespace, + now, optional_text, parsed_status, projection_facts, rejection, run_filter_spec, run_records, + run_rows, storage_error, stored_event, text_of, utc, utc_time, +}; +use crate::bindings::_helpers::execution_history::statements::{ + column, execute, execute_plain, execute_query, fetch_all, fetch_one, tuple_of, +}; + +const PSYCOPG_MODULE: &str = "psycopg"; +const PSYCOPG_ROWS_MODULE: &str = "psycopg.rows"; +const CLOSED: &str = "PostgreSQL execution history storage is closed"; +const MISSING_DRIVER: &str = "PostgreSQL execution history requires psycopg 3; install sqlbuild[postgres] or 'psycopg[binary]>=3.2'"; +const CONNECT_FAILED: &str = "cannot connect to PostgreSQL execution history; verify the secret-resolved DSN, network access, and sqlbuild[postgres] installation"; +const READ_FAILED: &str = "PostgreSQL execution history read failed"; +const METADATA_INVALID: &str = "PostgreSQL execution history schema metadata is invalid"; +const NAMESPACE_MISSING: &str = "PostgreSQL execution history cursor namespace is missing"; +const STORAGE_ERROR: &str = "ExecutionHistoryStorageError"; + +/// The `psycopg.Error` class. +fn driver_error_class(py: Python<'_>) -> PyResult> { + imported(py, PSYCOPG_MODULE, "Error") +} + +fn driver_failure(py: Python<'_>, error: &PyErr) -> PyResult { + Ok(error.is_instance(py, &driver_error_class(py)?)) +} + +/// Serialization failures and deadlocks retry while attempts remain. +fn retries_transaction(sqlstate: Option<&str>, attempts: i64, max_retries: i64) -> bool { + sqlstate.is_some_and(|state| POSTGRES_RETRYABLE_SQLSTATES.contains(&state)) + && attempts < max_retries +} + +/// `value` when it is an `int` that is not a `bool`. +fn storage_position(value: &Bound<'_, PyAny>, message: &str) -> PyResult { + exact_int(value).ok_or_else(|| storage_error(value.py(), message)) +} + +/// `value.astimezone(UTC)` when it is a `datetime`. +fn utc_datetime<'py>(value: &Bound<'py, PyAny>, message: &str) -> PyResult> { + let py: Python<'py> = value.py(); + if !value.is_instance_of::() { + return Err(storage_error(py, message)); + } + value.call_method1("astimezone", (utc(py)?,)) +} + +fn optional_utc(value: &Bound<'_, PyAny>) -> PyResult> { + if !value.is_instance_of::() { + return Ok(None); + } + utc_time(value).map(Some) +} + +fn text_column(row: &Bound<'_, PyAny>, name: &str) -> PyResult { + column(row, name)?.str()?.extract() +} + +/// One `sqlbuild_run_projection` row as a projected run. +fn stored_run(row: &Bound<'_, PyAny>) -> PyResult { + let py: Python<'_> = row.py(); + let created_at = utc_datetime( + &column(row, "created_at")?, + "PostgreSQL run timestamp is invalid", + )?; + let last_storage_order: i64 = storage_position( + &column(row, "last_event_storage_id")?, + "PostgreSQL run storage position is invalid", + )?; + let optional_str = |name: &str| -> PyResult> { + let value = column(row, name)?; + if value.is_none() { + return Ok(None); + } + Ok(Some(value.str()?.extract()?)) + }; + Ok(RunRow { + run_id: text_column(row, "run_id")?, + invocation_id: text_column(row, "invocation_id")?, + created_at: utc_time(&created_at)?, + status: parsed_status(py, &text_column(row, "status")?)?, + is_complete: column(row, "is_complete")?.is_truthy()?, + last_event_cursor: text_column(row, "last_event_cursor")?, + last_storage_order, + command: optional_str("command")?, + target: optional_str("target")?, + environment: optional_str("environment")?, + started_at: optional_utc(&column(row, "started_at")?)?, + ended_at: optional_utc(&column(row, "ended_at")?)?, + }) +} + +/// Deployed PostgreSQL implementation of event-log and run-projection contracts. +#[pyclass( + module = "sqlbuild._native", + name = "PostgresExecutionHistory", + frozen, + subclass, + dict +)] +pub(crate) struct PostgresExecutionHistory { + statement_timeout_ms: OnceLock, + max_transaction_retries: OnceLock, + closed: AtomicBool, + project_calls: AtomicI64, +} + +impl PostgresExecutionHistory { + fn ensure_open(&self, py: Python<'_>) -> PyResult<()> { + if self.closed.load(Ordering::SeqCst) { + return Err(storage_error(py, CLOSED)); + } + Ok(()) + } + + /// Run `action` in a transaction with the statement timeout, retrying serialization failures. + fn transactional<'py, T>( + &self, + slf: &Bound<'py, Self>, + action: impl Fn(&Bound<'py, PyAny>) -> PyResult, + failure_message: &str, + ) -> PyResult { + let py: Python<'py> = slf.py(); + let connection: Bound<'py, PyAny> = slf.getattr("_connection")?; + let timeout: String = self + .statement_timeout_ms + .get() + .copied() + .unwrap_or(POSTGRES_STATEMENT_TIMEOUT_MS) + .to_string(); + let max_retries: i64 = self + .max_transaction_retries + .get() + .copied() + .unwrap_or(POSTGRES_TRANSACTION_RETRIES); + let mut attempts: i64 = 0; + loop { + let result: PyResult = within(&connection.call_method0("transaction")?, |_| { + within(&connection.call_method0("cursor")?, |cursor| { + execute( + cursor, + POSTGRES_TRANSACTION_TIMEOUT, + &[SqlParam::Text(timeout.clone())], + )?; + action(cursor) + }) + }); + let error: PyErr = match result { + Ok(value) => return Ok(value), + Err(error) => error, + }; + if history_error_of_kind(py, &error, STORAGE_ERROR)? || !driver_failure(py, &error)? { + return Err(error); + } + let sqlstate: Option = text_of(&error.value(py).getattr("sqlstate")?)?; + if retries_transaction(sqlstate.as_deref(), attempts, max_retries) { + attempts += 1; + continue; + } + return Err(storage_error(py, failure_message)); + } + } + + /// `query` through a fresh cursor, with driver failures reported as read failures. + fn read_rows<'py>(slf: &Bound<'py, Self>, query: &Query) -> PyResult>> { + let py: Python<'py> = slf.py(); + let connection: Bound<'py, PyAny> = slf.getattr("_connection")?; + let rows = within(&connection.call_method0("cursor")?, |cursor| { + execute_query(cursor, query)?; + fetch_all(cursor) + }); + match rows { + Err(error) if driver_failure(py, &error)? => Err(storage_error(py, READ_FAILED)), + rows => rows, + } + } + + fn with_cursor<'py, T>( + slf: &Bound<'py, Self>, + cursor: Option<&Bound<'py, PyAny>>, + body: impl FnOnce(&Bound<'py, PyAny>) -> PyResult, + ) -> PyResult { + if let Some(cursor) = cursor { + return body(cursor); + } + let owned: Bound<'py, PyAny> = slf.getattr("_connection")?.call_method0("cursor")?; + let result: PyResult = body(&owned); + let closed = owned.call_method0("close"); + let value: T = result?; + closed?; + Ok(value) + } + + fn schema_version_in<'py>( + slf: &Bound<'py, Self>, + cursor: Option<&Bound<'py, PyAny>>, + ) -> PyResult { + let py: Python<'py> = slf.py(); + Self::with_cursor(slf, cursor, |active| { + let read = (|| -> PyResult { + execute_plain(active, POSTGRES_MIGRATIONS_EXIST)?; + let Some(exists) = fetch_one(active)? else { + return Ok(0); + }; + if exists.get_item("table_name")?.is_none() { + return Ok(0); + } + execute_plain(active, POSTGRES_SCHEMA_VERSION)?; + let row = fetch_one(active)?; + let value = row.map(|row| row.get_item("schema_version")).transpose()?; + value + .as_ref() + .and_then(exact_int) + .ok_or_else(|| storage_error(py, METADATA_INVALID)) + })(); + match read { + Err(error) if history_error_of_kind(py, &error, STORAGE_ERROR)? => Err(error), + Err(error) if error.is_instance_of::(py) => { + Err(storage_error(py, METADATA_INVALID)) + } + read => read, + } + }) + } + + fn namespace_in<'py>( + slf: &Bound<'py, Self>, + cursor: Option<&Bound<'py, PyAny>>, + ) -> PyResult { + let py: Python<'py> = slf.py(); + Self::with_cursor(slf, cursor, |active| { + execute_plain(active, POSTGRES_NAMESPACE)?; + let row = fetch_one(active)?; + let namespace = match row { + Some(row) => text_of(&row.get_item("storage_namespace")?)?, + None => None, + }; + namespace.ok_or_else(|| storage_error(py, NAMESPACE_MISSING)) + }) + } + + fn stored_event<'py>(row: &Bound<'py, PyAny>, namespace: &str) -> PyResult> { + let py: Python<'py> = row.py(); + let storage_id: i64 = storage_position( + &column(row, "storage_id")?, + "PostgreSQL event storage position is invalid", + )?; + let received_at = utc_datetime( + &column(row, "received_at")?, + "PostgreSQL event timestamp is invalid", + )?; + let payload: Bound<'py, PyAny> = column(row, "payload_json")?; + if text_of(&payload)?.is_none() { + return Err(storage_error( + py, + "PostgreSQL canonical event content is invalid", + )); + } + stored_event( + py, + ( + storage_id, + event_cursor(Dialect::Postgres, namespace, storage_id), + ), + received_at, + &payload, + ) + } + + fn stored_events<'py>( + slf: &Bound<'py, Self>, + rows: &[Bound<'py, PyAny>], + ) -> PyResult>> { + if rows.is_empty() { + return Ok(Vec::new()); + } + let namespace: String = Self::namespace_in(slf, None)?; + rows.iter() + .map(|row| Self::stored_event(row, &namespace)) + .collect() + } + + fn read_all_runs(cursor: &Bound<'_, PyAny>) -> PyResult> { + execute_plain(cursor, POSTGRES_ALL_RUNS)?; + fetch_all(cursor)?.iter().map(stored_run).collect() + } + + fn lock_history(cursor: &Bound<'_, PyAny>) -> PyResult<()> { + execute( + cursor, + POSTGRES_ADVISORY_LOCK, + &[SqlParam::Integer(POSTGRES_HISTORY_LOCK_ID)], + )?; + Ok(()) + } + + fn append_via<'py>( + slf: &Bound<'py, Self>, + cursor: &Bound<'py, PyAny>, + events: &Bound<'py, PyAny>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + slf.getattr("_append_events")?.call( + (), + Some(&keywords( + py, + &[("cursor", cursor.clone()), ("events", events.clone())], + )?), + ) + } + + fn replace_via<'py>( + slf: &Bound<'py, Self>, + cursor: &Bound<'py, PyAny>, + projected: &Bound<'py, PyAny>, + ) -> PyResult<()> { + let py: Python<'py> = slf.py(); + slf.getattr("_replace_projection")?.call( + (), + Some(&keywords( + py, + &[("cursor", cursor.clone()), ("projected", projected.clone())], + )?), + )?; + Ok(()) + } + + fn migrate(&self, slf: &Bound<'_, Self>) -> PyResult<()> { + let py: Python<'_> = slf.py(); + self.transactional( + slf, + |cursor| { + execute( + cursor, + POSTGRES_ADVISORY_LOCK, + &[SqlParam::Integer(POSTGRES_MIGRATION_LOCK_ID)], + )?; + let current: i64 = Self::schema_version_in(slf, Some(cursor))?; + if !checked(py, schema_migration_needed(Dialect::Postgres, current))? { + return Ok(()); + } + for statement in POSTGRES_SCHEMA { + execute_plain(cursor, statement)?; + } + execute( + cursor, + POSTGRES_RECORD_MIGRATION, + &[ + SqlParam::Integer(SCHEMA_VERSION), + SqlParam::Text(new_namespace(py)?), + ], + )?; + Ok(()) + }, + "PostgreSQL execution history schema migration failed", + ) + } + + fn run_cursor_key( + slf: &Bound<'_, Self>, + cursor: Option<&Bound<'_, PyAny>>, + ) -> PyResult> { + let py: Python<'_> = slf.py(); + let Some(cursor) = cursor.filter(|cursor| !cursor.is_none()) else { + return Ok(None); + }; + let text: String = text_of(cursor)?.unwrap_or_default(); + let namespace: String = Self::namespace_in(slf, None)?; + let decoded = checked(py, decoded_run_cursor(Dialect::Postgres, &text, &namespace))?; + let invalid = || rejection(py, HistoryErrorKind::InvalidCursor, RUN_CURSOR_TIMESTAMP); + let parsed = imported(py, "datetime", "datetime")? + .call_method1("fromisoformat", (decoded.created_at_text.as_str(),)); + let parsed = match parsed { + Ok(parsed) => parsed, + Err(error) if error.is_instance_of::(py) => { + return Err(invalid()); + } + Err(error) => return Err(error), + }; + if parsed.call_method0("utcoffset")?.is_none() { + return Err(invalid()); + } + let created_at = utc_time(&parsed)?; + let rows = Self::read_rows( + slf, + &Query { + sql: POSTGRES_RUN_EXISTS.to_owned(), + params: vec![ + SqlParam::Timestamp(created_at), + SqlParam::Text(decoded.run_id.clone()), + ], + }, + )?; + if rows.is_empty() { + return Err(rejection( + py, + HistoryErrorKind::InvalidCursor, + UNKNOWN_RUN_CURSOR, + )); + } + Ok(Some(RunCursorKey { + created_at, + created_at_text: decoded.created_at_text, + run_id: decoded.run_id, + })) + } +} + +#[pymethods] +impl PostgresExecutionHistory { + #[new] + #[pyo3(signature = (*_args, **_kwargs))] + fn new(_args: &Bound<'_, PyTuple>, _kwargs: Option<&Bound<'_, PyDict>>) -> Self { + Self { + statement_timeout_ms: OnceLock::new(), + max_transaction_retries: OnceLock::new(), + closed: AtomicBool::new(true), + project_calls: AtomicI64::new(0), + } + } + + #[pyo3(signature = (dsn, **options))] + fn __init__( + slf: &Bound<'_, Self>, + dsn: &Bound<'_, PyAny>, + options: Option>, + ) -> PyResult<()> { + let py: Python<'_> = slf.py(); + let this: &Self = slf.get(); + let options = KeywordArguments::read( + py, + "__init__", + options, + ( + &[], + &[ + "connect_timeout_seconds", + "statement_timeout_ms", + "max_transaction_retries", + ], + ), + )?; + let setting = |name: &str, default: i64| -> PyResult> { + if options.passed(name)? { + return Ok(exact_int(&options.value(name)?)); + } + Ok(Some(default)) + }; + let settings = PostgresSettings { + dsn: text_of(dsn)?, + connect_timeout_seconds: setting( + "connect_timeout_seconds", + POSTGRES_CONNECT_TIMEOUT_SECONDS, + )?, + statement_timeout_ms: setting("statement_timeout_ms", POSTGRES_STATEMENT_TIMEOUT_MS)?, + max_transaction_retries: setting( + "max_transaction_retries", + POSTGRES_TRANSACTION_RETRIES, + )?, + }; + checked(py, validate_postgres_settings(&settings))?; + let (psycopg, dict_row) = match ( + py.import(PSYCOPG_MODULE), + imported(py, PSYCOPG_ROWS_MODULE, "dict_row"), + ) { + (Ok(psycopg), Ok(dict_row)) => (psycopg, dict_row), + (Err(error), _) | (_, Err(error)) + if error.is_instance_of::(py) => + { + return Err(storage_error(py, MISSING_DRIVER)); + } + (Err(error), _) | (_, Err(error)) => return Err(error), + }; + let statement_timeout: i64 = settings + .statement_timeout_ms + .unwrap_or(POSTGRES_STATEMENT_TIMEOUT_MS); + let _ = this.max_transaction_retries.set( + settings + .max_transaction_retries + .unwrap_or(POSTGRES_TRANSACTION_RETRIES), + ); + this.closed.store(false, Ordering::SeqCst); + let mut connection: Option> = None; + let opened: PyResult<()> = (|| { + let opened: Bound<'_, PyAny> = psycopg.getattr("connect")?.call( + (dsn,), + Some(&keywords( + py, + &[ + ( + "connect_timeout", + settings.connect_timeout_seconds.into_bound_py_any(py)?, + ), + ("autocommit", PyBool::new(py, true).to_owned().into_any()), + ("row_factory", dict_row.clone()), + ], + )?), + )?; + connection = Some(opened.clone()); + slf.setattr("_connection", &opened)?; + within(&opened.call_method0("cursor")?, |cursor| { + execute( + cursor, + POSTGRES_SESSION_TIMEOUT, + &[SqlParam::Text(statement_timeout.to_string())], + ) + })?; + let _ = this.statement_timeout_ms.set(statement_timeout); + this.migrate(slf)?; + slf.call_method0("reconcile")?; + Ok(()) + })(); + let Err(error) = opened else { + return Ok(()); + }; + if !error.is_instance_of::(py) { + return Err(error); + } + if let Some(connection) = connection { + let _ = connection.call_method0("close"); + } + this.closed.store(true, Ordering::SeqCst); + if history_error_of_kind(py, &error, STORAGE_ERROR)? { + return Err(error); + } + Err(storage_error(py, CONNECT_FAILED)) + } + + fn __repr__(slf: &Bound<'_, Self>) -> PyResult { + let state: &str = if slf.get().closed.load(Ordering::SeqCst) { + "closed" + } else { + "open" + }; + let name: String = slf.get_type().getattr("__name__")?.extract()?; + Ok(format!("{name}(state='{state}')")) + } + + /// The number of incremental projection calls. + #[getter] + fn project_calls(&self) -> i64 { + self.project_calls.load(Ordering::SeqCst) + } + + fn append_event<'py>( + slf: &Bound<'py, Self>, + event: &Bound<'py, PyAny>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + slf.get().ensure_open(py)?; + slf.call_method1("append_events", (tuple_of(py, vec![event.clone()])?,))? + .get_item(0) + } + + fn append_events<'py>( + slf: &Bound<'py, Self>, + events: &Bound<'py, PyAny>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + let this: &Self = slf.get(); + this.ensure_open(py)?; + let pending: Bound<'py, PyTuple> = + PyTuple::new(py, events.try_iter()?.collect::>>()?)?; + this.transactional( + slf, + |cursor| { + Self::lock_history(cursor)?; + Self::append_via(slf, cursor, pending.as_any()) + }, + "PostgreSQL event append failed", + ) + } + + #[pyo3(signature = (*, event_filter, after_cursor = None, limit = None))] + #[pyo3(name = "get_events")] + fn stored_events_page<'py>( + slf: &Bound<'py, Self>, + event_filter: &Bound<'py, PyAny>, + after_cursor: Option<&Bound<'py, PyAny>>, + limit: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + slf.get().ensure_open(py)?; + let limit: Option = limit.map_or(Some(DEFAULT_PAGE_LIMIT), exact_int); + checked(py, validate_page_limit(limit))?; + let limit: i64 = limit.unwrap_or(DEFAULT_PAGE_LIMIT); + let after: i64 = match after_cursor.filter(|cursor| !cursor.is_none()) { + None => 0, + Some(cursor) => { + let text: String = text_of(cursor)?.unwrap_or_default(); + let namespace: String = Self::namespace_in(slf, None)?; + let position: i64 = checked( + py, + event_cursor_position(Dialect::Postgres, Some(&text), &namespace), + )?; + let rows = Self::read_rows( + slf, + &Query { + sql: POSTGRES_EVENT_EXISTS.to_owned(), + params: vec![SqlParam::Integer(position)], + }, + )?; + if rows.is_empty() { + return Err(rejection( + py, + HistoryErrorKind::InvalidCursor, + UNKNOWN_EVENT_CURSOR, + )); + } + position + } + }; + let spec = event_filter_spec(event_filter)?; + let rows = Self::read_rows( + slf, + &event_page_query(Dialect::Postgres, after, &spec, limit), + )?; + let has_more: bool = rows.len() as i64 > limit; + let page: Vec> = rows.into_iter().take(limit as usize).collect(); + let records: Vec> = Self::stored_events(slf, &page)?; + let next_cursor: Bound<'py, PyAny> = match records.last() { + Some(record) => record.getattr("cursor")?, + None => py.None().into_bound(py), + }; + call_with( + &imported(py, MODELS_MODULE, "EventPage")?, + &[ + ("records", tuple_of(py, records)?.into_any()), + ("next_cursor", next_cursor), + ("has_more", has_more.into_bound_py_any(py)?), + ], + ) + } + + #[pyo3(name = "get_run")] + fn run_record<'py>( + slf: &Bound<'py, Self>, + run_id: &Bound<'py, PyAny>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + slf.get().ensure_open(py)?; + let run_id: Option = optional_text(run_id)?; + let params: Vec = vec![run_id.map_or(SqlParam::Null, SqlParam::Text)]; + let rows = Self::read_rows( + slf, + &Query { + sql: POSTGRES_RUN_BY_ID.to_owned(), + params, + }, + )?; + match rows.first() { + Some(row) => crate::bindings::_helpers::execution_history::conversion::run_record( + py, + &stored_run(row)?, + ), + None => Ok(py.None().into_bound(py)), + } + } + + #[pyo3(signature = (*, run_filter, after_cursor = None, limit = None))] + #[pyo3(name = "get_runs")] + fn run_page<'py>( + slf: &Bound<'py, Self>, + run_filter: &Bound<'py, PyAny>, + after_cursor: Option<&Bound<'py, PyAny>>, + limit: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + slf.get().ensure_open(py)?; + let limit: Option = limit.map_or(Some(DEFAULT_PAGE_LIMIT), exact_int); + checked(py, validate_page_limit(limit))?; + let limit: i64 = limit.unwrap_or(DEFAULT_PAGE_LIMIT); + let after = Self::run_cursor_key(slf, after_cursor)?; + let spec = run_filter_spec(run_filter)?; + let rows = Self::read_rows( + slf, + &run_page_query(Dialect::Postgres, after.as_ref(), &spec, limit), + )?; + let runs: Vec = rows + .iter() + .take(limit as usize) + .map(stored_run) + .collect::>()?; + let next_cursor: Bound<'py, PyAny> = match runs.last() { + Some(run) => { + let namespace: String = Self::namespace_in(slf, None)?; + checked(py, run_cursor(Dialect::Postgres, &namespace, run))? + .into_bound_py_any(py)? + } + None => py.None().into_bound(py), + }; + call_with( + &imported(py, MODELS_MODULE, "RunPage")?, + &[ + ("records", run_records(py, &runs)?.into_any()), + ("next_cursor", next_cursor), + ( + "has_more", + (rows.len() as i64 > limit).into_bound_py_any(py)?, + ), + ], + ) + } + + fn project<'py>( + slf: &Bound<'py, Self>, + stored_events: &Bound<'py, PyAny>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + let this: &Self = slf.get(); + this.ensure_open(py)?; + let pending: Bound<'py, PyTuple> = + PyTuple::new(py, stored_events.try_iter()?.collect::>>()?)?; + this.project_calls.fetch_add(1, Ordering::SeqCst); + this.transactional( + slf, + |cursor| { + Self::lock_history(cursor)?; + let facts = projection_facts(pending.as_any())?; + let projected = checked(py, project_runs(facts, Self::read_all_runs(cursor)?))?; + let records: Bound<'py, PyTuple> = run_records(py, &projected)?; + Self::replace_via(slf, cursor, records.as_any())?; + Ok(records) + }, + "PostgreSQL run projection publication failed", + ) + } + + fn rebuild_from_events<'py>( + slf: &Bound<'py, Self>, + stored_events: &Bound<'py, PyAny>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + let this: &Self = slf.get(); + this.ensure_open(py)?; + let pending: Bound<'py, PyTuple> = + PyTuple::new(py, stored_events.try_iter()?.collect::>>()?)?; + this.transactional( + slf, + |cursor| { + Self::lock_history(cursor)?; + let projected = checked( + py, + project_runs(projection_facts(pending.as_any())?, Vec::new()), + )?; + let records: Bound<'py, PyTuple> = run_records(py, &projected)?; + Self::replace_via(slf, cursor, records.as_any())?; + Ok(records) + }, + "PostgreSQL run projection rebuild failed", + ) + } + + /// Atomically append immutable facts and publish their run projection. + fn append_and_project<'py>( + slf: &Bound<'py, Self>, + events: &Bound<'py, PyAny>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + let this: &Self = slf.get(); + this.ensure_open(py)?; + let pending: Bound<'py, PyTuple> = + PyTuple::new(py, events.try_iter()?.collect::>>()?)?; + this.transactional( + slf, + |cursor| { + Self::lock_history(cursor)?; + let stored: Bound<'py, PyAny> = Self::append_via(slf, cursor, pending.as_any())?; + let projected = checked( + py, + project_runs(projection_facts(&stored)?, Self::read_all_runs(cursor)?), + )?; + let records: Bound<'py, PyTuple> = run_records(py, &projected)?; + Self::replace_via(slf, cursor, records.as_any())?; + Ok(stored) + }, + "PostgreSQL atomic append and projection failed", + ) + } + + /// Rebuild the disposable run projection from all durable event facts. + fn reconcile<'py>(slf: &Bound<'py, Self>) -> PyResult> { + let py: Python<'py> = slf.py(); + let this: &Self = slf.get(); + this.ensure_open(py)?; + this.transactional( + slf, + |cursor| { + Self::lock_history(cursor)?; + let stored: Bound<'py, PyAny> = slf + .getattr("_read_all_events")? + .call((), Some(&keywords(py, &[("cursor", cursor.clone())])?))?; + let projected = checked(py, project_runs(projection_facts(&stored)?, Vec::new()))?; + let records: Bound<'py, PyTuple> = run_records(py, &projected)?; + Self::replace_via(slf, cursor, records.as_any())?; + Ok(records) + }, + "PostgreSQL reconciliation failed", + ) + } + + /// Verify connectivity and the supported schema revision. + fn check_health(slf: &Bound<'_, Self>) -> PyResult { + let py: Python<'_> = slf.py(); + slf.get().ensure_open(py)?; + let healthy = (|| -> PyResult { + let connection: Bound<'_, PyAny> = slf.getattr("_connection")?; + let row = within(&connection.call_method0("cursor")?, |cursor| { + execute_plain(cursor, POSTGRES_PING)?; + fetch_one(cursor) + })?; + let version: i64 = slf.call_method0("get_schema_version")?.extract()?; + Ok(row.is_some() && version == SCHEMA_VERSION) + })(); + match healthy { + Err(error) if driver_failure(py, &error)? => Ok(false), + healthy => healthy, + } + } + + #[pyo3(name = "get_schema_version")] + fn schema_version(slf: &Bound<'_, Self>) -> PyResult { + slf.get().ensure_open(slf.py())?; + Self::schema_version_in(slf, None) + } + + #[pyo3(signature = (*, target_version = None))] + fn upgrade_schema(slf: &Bound<'_, Self>, target_version: Option) -> PyResult { + let py: Python<'_> = slf.py(); + let this: &Self = slf.get(); + this.ensure_open(py)?; + let target: i64 = target_version.unwrap_or(SCHEMA_VERSION); + checked(py, validate_schema_target(Dialect::Postgres, target))?; + this.migrate(slf)?; + Ok(target) + } + + fn close(slf: &Bound<'_, Self>) -> PyResult<()> { + let this: &Self = slf.get(); + if this.closed.load(Ordering::SeqCst) { + return Ok(()); + } + slf.getattr("_connection")?.call_method0("close")?; + this.closed.store(true, Ordering::SeqCst); + Ok(()) + } + + fn dispose(slf: &Bound<'_, Self>) -> PyResult<()> { + slf.call_method0("close")?; + Ok(()) + } + + fn __enter__<'py>(slf: &Bound<'py, Self>) -> PyResult> { + slf.get().ensure_open(slf.py())?; + Ok(slf.clone()) + } + + #[pyo3(signature = (*_exc_info))] + fn __exit__(slf: &Bound<'_, Self>, _exc_info: &Bound<'_, PyTuple>) -> PyResult> { + slf.call_method0("close")?; + Ok(None) + } + + #[pyo3(signature = (*, cursor, events))] + fn _append_events<'py>( + slf: &Bound<'py, Self>, + cursor: &Bound<'py, PyAny>, + events: &Bound<'py, PyAny>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + let prepared: Vec = prepared_batch(events)?; + let mut stored: Vec> = Vec::with_capacity(prepared.len()); + let mut namespace: Option = None; + let received_at = now(py)?; + for event in &prepared { + execute_query( + cursor, + &event_insert_query(Dialect::Postgres, event, &received_at), + )?; + let mut row = fetch_one(cursor)?; + if row.is_none() { + execute( + cursor, + POSTGRES_EVENT_BY_ID, + &[SqlParam::Text(event.event_id.clone())], + )?; + row = fetch_one(cursor)?; + } + let Some(row) = row else { + return Err(storage_error( + py, + "PostgreSQL event append did not produce a durable fact", + )); + }; + let digest: Option = text_of(&column(&row, "content_digest")?)?; + let payload: Option = text_of(&column(&row, "payload_json")?)?; + checked( + py, + validate_stored_content(event, (digest.as_deref(), payload.as_deref())), + )?; + if namespace.is_none() { + namespace = Some(Self::namespace_in(slf, Some(cursor))?); + } + stored.push(Self::stored_event( + &row, + namespace.as_deref().unwrap_or_default(), + )?); + } + tuple_of(py, stored) + } + + #[pyo3(signature = (*, cursor, projected))] + fn _replace_projection( + _slf: &Bound<'_, Self>, + cursor: &Bound<'_, PyAny>, + projected: &Bound<'_, PyAny>, + ) -> PyResult<()> { + execute_plain(cursor, POSTGRES_DELETE_RUNS)?; + for run in run_rows(projected)? { + execute_query(cursor, &run_insert_query(Dialect::Postgres, &run))?; + } + Ok(()) + } + + #[pyo3(signature = (*, cursor))] + fn _read_all_events<'py>( + slf: &Bound<'py, Self>, + cursor: &Bound<'py, PyAny>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + execute_plain(cursor, POSTGRES_ALL_EVENTS)?; + let rows: Vec> = fetch_all(cursor)?; + let stored: Vec> = if rows.is_empty() { + Vec::new() + } else { + let namespace: String = Self::namespace_in(slf, Some(cursor))?; + rows.iter() + .map(|row| Self::stored_event(row, &namespace)) + .collect::>()? + }; + tuple_of(py, stored) + } +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/execution_history/registration.rs b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/registration.rs new file mode 100644 index 0000000000..52d85a14b2 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/registration.rs @@ -0,0 +1,13 @@ +//! Register native execution history storage and entry points. + +use pyo3::prelude::{Bound, PyModule, PyModuleMethods, PyResult}; + +use crate::bindings::_helpers::execution_history::functions; +use crate::bindings::_helpers::execution_history::postgres::PostgresExecutionHistory; +use crate::bindings::_helpers::execution_history::sqlite::SqliteExecutionHistory; + +pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { + module.add_class::()?; + module.add_class::()?; + functions::register(module) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/execution_history/sqlite.rs b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/sqlite.rs new file mode 100644 index 0000000000..9148d2a76d --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/sqlite.rs @@ -0,0 +1,921 @@ +//! `SQLiteExecutionHistory`: project-local SQLite event-log and run-projection storage. + +use std::fs::{DirBuilder, Permissions}; +use std::io::ErrorKind; +use std::os::unix::fs::{DirBuilderExt, PermissionsExt}; +use std::path::Path; +use std::sync::OnceLock; +use std::sync::atomic::{AtomicBool, AtomicI64, Ordering}; + +use pyo3::IntoPyObjectExt; +use pyo3::prelude::{Bound, Py, PyAny, PyAnyMethods, PyErr, PyResult, Python, pyclass, pymethods}; +use pyo3::types::{PyDict, PyTuple}; +use sqlbuild_state::execution_history::constants::{ + DEFAULT_PAGE_LIMIT, MEMORY_PATH, RUN_CURSOR_TIMESTAMP, SCHEMA_VERSION, SQLITE_ALL_EVENTS, + SQLITE_ALL_RUNS, SQLITE_BEGIN, SQLITE_BUSY_TIMEOUT_MS, SQLITE_COMMIT, SQLITE_DELETE_RUNS, + SQLITE_DIRECTORY_MODE, SQLITE_EVENT_BY_ID, SQLITE_EVENT_BY_STORAGE_ID, SQLITE_EVENT_EXISTS, + SQLITE_FILE_MODE, SQLITE_FOREIGN_KEYS, SQLITE_HEALTHY_CHECK_RESULT, SQLITE_INSERT_NAMESPACE, + SQLITE_INSERT_SCHEMA_VERSION, SQLITE_JOURNAL_MODE, SQLITE_METADATA_EXISTS, + SQLITE_METADATA_TABLE, SQLITE_METADATA_VALUE, SQLITE_NAMESPACE_KEY, SQLITE_QUICK_CHECK, + SQLITE_ROLLBACK, SQLITE_RUN_BY_ID, SQLITE_RUN_EXISTS, SQLITE_RUN_INSERT, SQLITE_SCHEMA, + SQLITE_SCHEMA_VERSION_KEY, UNKNOWN_EVENT_CURSOR, UNKNOWN_RUN_CURSOR, +}; +use sqlbuild_state::execution_history::errors::HistoryErrorKind; +use sqlbuild_state::execution_history::main::decoded_run_cursor::decoded_run_cursor; +use sqlbuild_state::execution_history::main::event_cursor::event_cursor; +use sqlbuild_state::execution_history::main::event_cursor_position::event_cursor_position; +use sqlbuild_state::execution_history::main::event_insert_query::event_insert_query; +use sqlbuild_state::execution_history::main::event_page_query::event_page_query; +use sqlbuild_state::execution_history::main::project_runs::project_runs; +use sqlbuild_state::execution_history::main::run_cursor::run_cursor; +use sqlbuild_state::execution_history::main::run_insert_query::run_insert_query; +use sqlbuild_state::execution_history::main::run_page_query::run_page_query; +use sqlbuild_state::execution_history::main::schema_migration_needed::schema_migration_needed; +use sqlbuild_state::execution_history::main::sqlite_database_path::sqlite_database_path; +use sqlbuild_state::execution_history::main::validate_page_limit::validate_page_limit; +use sqlbuild_state::execution_history::main::validate_schema_target::validate_schema_target; +use sqlbuild_state::execution_history::main::validate_stored_content::validate_stored_content; +use sqlbuild_state::execution_history::models::{ + Dialect, PreparedEvent, RunCursorKey, RunRow, SqlParam, SqliteLocation, +}; +use sqlbuild_state::execution_history::types::UtcTime; + +use crate::bindings::_helpers::compile_driver::python_calls::imported; +use crate::bindings::_helpers::execution_history::appends::prepared_batch; +use crate::bindings::_helpers::execution_history::conversion::{ + checked, event_filter_spec, exact_int, history_error_of_kind, new_namespace, now, + optional_text, parsed_iso, parsed_status, projection_facts, python_params, rejection, + run_filter_spec, run_records, run_rows, storage_error, storage_error_from, stored_event, + text_of, utc_time, +}; +use crate::bindings::_helpers::execution_history::statements::{ + column, execute, execute_plain, execute_query, fetch_one, locked, tuple_of, +}; + +const SQLITE_MODULE: &str = "sqlite3"; +const PATHLIB_MODULE: &str = "pathlib"; +const THREADING_MODULE: &str = "threading"; +const MILLIS_PER_SECOND: f64 = 1000.0; +const CLOSED: &str = "SQLite execution history storage is closed"; +const NOT_AWARE: &str = "SQLite history timestamp is not timezone-aware"; +const INTEGRITY_CONFLICT: &str = "IntegrityConflictError"; +const STORAGE_ERROR: &str = "ExecutionHistoryStorageError"; + +/// `Path(path).parent.mkdir(mode=0o700, parents=True, exist_ok=True)` for a file-backed history. +fn prepare_history_directory(path: &str) -> std::io::Result<()> { + if path == MEMORY_PATH { + return Ok(()); + } + let Some(parent) = Path::new(path).parent() else { + return Ok(()); + }; + if parent.as_os_str().is_empty() || parent.is_dir() { + return Ok(()); + } + if let Some(grandparent) = parent.parent() + && !grandparent.as_os_str().is_empty() + { + std::fs::create_dir_all(grandparent)?; + } + match DirBuilder::new().mode(SQLITE_DIRECTORY_MODE).create(parent) { + Err(error) if error.kind() == ErrorKind::AlreadyExists && parent.is_dir() => Ok(()), + result => result, + } +} + +/// `os.chmod(path, 0o600)` for a file-backed history. +fn restrict_history_file(path: &str) -> std::io::Result<()> { + if path == MEMORY_PATH { + return Ok(()); + } + std::fs::set_permissions(path, Permissions::from_mode(SQLITE_FILE_MODE)) +} + +/// Whether `error` is a `sqlite3.Error`. +fn sqlite_failure(py: Python<'_>, error: &PyErr) -> PyResult { + Ok(error.is_instance(py, &imported(py, SQLITE_MODULE, "Error")?)) +} + +/// `result`, with a `sqlite3.Error` raised as a storage error `message` from it. +fn wrapped(py: Python<'_>, result: PyResult, message: &str) -> PyResult { + match result { + Err(error) if sqlite_failure(py, &error)? => Err(storage_error_from(py, message, error)), + result => result, + } +} + +/// A stored ISO timestamp as an aware UTC `datetime`. +fn parsed_timestamp<'py>(value: &Bound<'py, PyAny>) -> PyResult> { + let py: Python<'py> = value.py(); + let text: String = value.extract()?; + let parsed: Bound<'py, PyAny> = imported(py, "datetime", "datetime")? + .call_method1("fromisoformat", (text.replace('Z', "+00:00"),))?; + if parsed.call_method0("utcoffset")?.is_none() { + return Err(storage_error(py, NOT_AWARE)); + } + parsed.call_method1("astimezone", (imported(py, "datetime", "UTC")?,)) +} + +fn optional_timestamp(value: &Bound<'_, PyAny>) -> PyResult> { + if value.is_none() { + return Ok(None); + } + utc_time(&parsed_timestamp(value)?).map(Some) +} + +/// One `run_projection` row as a projected run. +fn stored_run(row: &Bound<'_, PyAny>) -> PyResult { + let py: Python<'_> = row.py(); + let status: String = column(row, "status")?.extract()?; + Ok(RunRow { + run_id: column(row, "run_id")?.extract()?, + invocation_id: column(row, "invocation_id")?.extract()?, + created_at: utc_time(&parsed_timestamp(&column(row, "created_at")?)?)?, + status: parsed_status(py, &status)?, + is_complete: column(row, "is_complete")?.is_truthy()?, + last_event_cursor: column(row, "last_event_cursor")?.extract()?, + last_storage_order: column(row, "last_storage_order")?.extract()?, + command: optional_text(&column(row, "command")?)?, + target: optional_text(&column(row, "target")?)?, + environment: optional_text(&column(row, "environment")?)?, + started_at: optional_timestamp(&column(row, "started_at")?)?, + ended_at: optional_timestamp(&column(row, "ended_at")?)?, + }) +} + +/// Project-local SQLite implementation of event-log and run-projection storage contracts. +#[pyclass( + module = "sqlbuild._native", + name = "SQLiteExecutionHistory", + frozen, + subclass, + dict +)] +pub(crate) struct SqliteExecutionHistory { + path: OnceLock, + lock: OnceLock>, + closed: AtomicBool, + project_calls: AtomicI64, +} + +impl SqliteExecutionHistory { + fn database_path(&self) -> &str { + self.path.get().map_or(MEMORY_PATH, String::as_str) + } + + fn with_lock<'py, T>( + &self, + py: Python<'py>, + body: impl FnOnce() -> PyResult, + ) -> PyResult { + let Some(lock) = self.lock.get() else { + return Err(storage_error(py, CLOSED)); + }; + locked(lock.bind(py), body) + } + + fn ensure_open(&self, py: Python<'_>) -> PyResult<()> { + if self.closed.load(Ordering::SeqCst) { + return Err(storage_error(py, CLOSED)); + } + Ok(()) + } + + /// Run `body` in `BEGIN IMMEDIATE`, committing on success and rolling back on failure. + fn transaction<'py, T>( + &self, + slf: &Bound<'py, Self>, + body: impl FnOnce(&Bound<'py, PyAny>) -> PyResult, + ) -> PyResult { + let py: Python<'py> = slf.py(); + self.with_lock(py, || { + self.ensure_open(py)?; + let connection: Bound<'py, PyAny> = slf.getattr("_connection")?; + execute_plain(&connection, SQLITE_BEGIN)?; + let value: T = match body(&connection) { + Ok(value) => value, + Err(error) => { + let _ = execute_plain(&connection, SQLITE_ROLLBACK); + return Err(error); + } + }; + if let Err(error) = execute_plain(&connection, SQLITE_COMMIT) { + if sqlite_failure(py, &error)? + && connection.getattr("in_transaction")?.is_truthy()? + { + let _ = execute_plain(&connection, SQLITE_ROLLBACK); + } + return Err(error); + } + Ok(value) + }) + } + + fn metadata_value<'py>( + connection: &Bound<'py, PyAny>, + key: &str, + ) -> PyResult>> { + let row = fetch_one(&execute( + connection, + SQLITE_METADATA_VALUE, + &[SqlParam::Text(key.to_owned())], + )?)?; + row.map(|row| row.get_item(0)).transpose() + } + + fn metadata_version(connection: &Bound<'_, PyAny>) -> PyResult { + let py: Python<'_> = connection.py(); + let exists = fetch_one(&execute( + connection, + SQLITE_METADATA_EXISTS, + &[SqlParam::Text(SQLITE_METADATA_TABLE.to_owned())], + )?)?; + if exists.is_none() { + return Ok(0); + } + let Some(value) = Self::metadata_value(connection, SQLITE_SCHEMA_VERSION_KEY)? else { + return Ok(0); + }; + let invalid = || storage_error(py, "SQLite history schema metadata is invalid"); + match text_of(&value)? { + Some(text) => text.trim().parse::().map_err(|_| invalid()), + None => exact_int(&value).ok_or_else(invalid), + } + } + + fn namespace(connection: &Bound<'_, PyAny>) -> PyResult { + let py: Python<'_> = connection.py(); + match Self::metadata_value(connection, SQLITE_NAMESPACE_KEY)? { + Some(value) => Ok(value.str()?.extract()?), + None => Err(storage_error( + py, + "SQLite history cursor namespace is missing", + )), + } + } + + fn stored_event<'py>(row: &Bound<'py, PyAny>, namespace: &str) -> PyResult> { + let py: Python<'py> = row.py(); + let storage_id: i64 = column(row, "storage_id")?.extract()?; + stored_event( + py, + ( + storage_id, + event_cursor(Dialect::Sqlite, namespace, storage_id), + ), + parsed_timestamp(&column(row, "received_at")?)?, + &column(row, "payload_json")?, + ) + } + + fn read_all_runs(connection: &Bound<'_, PyAny>) -> PyResult> { + execute_plain(connection, SQLITE_ALL_RUNS)? + .try_iter()? + .map(|row| row.and_then(|row| stored_run(&row))) + .collect() + } + + fn replace_projection(connection: &Bound<'_, PyAny>, projected: &[RunRow]) -> PyResult<()> { + let py: Python<'_> = connection.py(); + execute_plain(connection, SQLITE_DELETE_RUNS)?; + let rows: Vec> = projected + .iter() + .map(|run| { + python_params(py, &run_insert_query(Dialect::Sqlite, run).params) + .map(Bound::into_any) + }) + .collect::>()?; + connection.call_method1("executemany", (SQLITE_RUN_INSERT, tuple_of(py, rows)?))?; + Ok(()) + } + + fn append_prepared<'py>( + connection: &Bound<'py, PyAny>, + prepared: &[PreparedEvent], + ) -> PyResult>> { + let py: Python<'py> = connection.py(); + let mut stored: Vec> = Vec::with_capacity(prepared.len()); + let mut namespace: Option = None; + for event in prepared { + let existing = fetch_one(&execute( + connection, + SQLITE_EVENT_BY_ID, + &[SqlParam::Text(event.event_id.clone())], + )?)?; + let row: Bound<'py, PyAny> = match existing { + Some(existing) => { + let digest: Option = text_of(&column(&existing, "content_digest")?)?; + let payload: Option = text_of(&column(&existing, "payload_json")?)?; + checked( + py, + validate_stored_content(event, (digest.as_deref(), payload.as_deref())), + )?; + existing + } + None => { + let query = event_insert_query(Dialect::Sqlite, event, &now(py)?); + let cursor: Bound<'py, PyAny> = execute_query(connection, &query)?; + let storage_id: Bound<'py, PyAny> = cursor.getattr("lastrowid")?; + connection + .call_method1("execute", (SQLITE_EVENT_BY_STORAGE_ID, (storage_id,)))? + .call_method0("fetchone")? + } + }; + if namespace.is_none() { + namespace = Some(Self::namespace(connection)?); + } + let namespace: &str = namespace.as_deref().unwrap_or_default(); + stored.push(Self::stored_event(&row, namespace)?); + } + Ok(stored) + } + + fn append_events_in<'py>( + connection: &Bound<'py, PyAny>, + events: &Bound<'py, PyAny>, + ) -> PyResult>> { + let prepared: Vec = prepared_batch(events)?; + Self::append_prepared(connection, &prepared) + } + + fn migrate_in(&self, slf: &Bound<'_, Self>) -> PyResult<()> { + let py: Python<'_> = slf.py(); + let connection: Bound<'_, PyAny> = slf.getattr("_connection")?; + let current: i64 = Self::metadata_version(&connection)?; + if !checked(py, schema_migration_needed(Dialect::Sqlite, current))? { + return Ok(()); + } + let namespace: String = new_namespace(py)?; + let migrated = self.transaction(slf, |connection| { + for statement in SQLITE_SCHEMA { + execute_plain(connection, statement)?; + } + execute( + connection, + SQLITE_INSERT_NAMESPACE, + &[ + SqlParam::Text(SQLITE_NAMESPACE_KEY.to_owned()), + SqlParam::Text(namespace.clone()), + ], + )?; + execute( + connection, + SQLITE_INSERT_SCHEMA_VERSION, + &[ + SqlParam::Text(SQLITE_SCHEMA_VERSION_KEY.to_owned()), + SqlParam::Text(SCHEMA_VERSION.to_string()), + ], + )?; + Ok(()) + }); + wrapped(py, migrated, "SQLite history schema migration failed") + } + + fn run_cursor_key( + connection: &Bound<'_, PyAny>, + cursor: &Bound<'_, PyAny>, + ) -> PyResult> { + let py: Python<'_> = connection.py(); + if cursor.is_none() { + return Ok(None); + } + let text: String = text_of(cursor)?.unwrap_or_default(); + let namespace: String = Self::namespace(connection)?; + let decoded = checked(py, decoded_run_cursor(Dialect::Sqlite, &text, &namespace))?; + let invalid_timestamp = + || rejection(py, HistoryErrorKind::InvalidCursor, RUN_CURSOR_TIMESTAMP); + let created_at = match parsed_iso(py, &decoded.created_at_text)? { + Some(parsed) if !parsed.call_method0("utcoffset")?.is_none() => utc_time(&parsed)?, + _ => return Err(invalid_timestamp()), + }; + let exists = fetch_one(&execute( + connection, + SQLITE_RUN_EXISTS, + &[ + SqlParam::Text(decoded.created_at_text.clone()), + SqlParam::Text(decoded.run_id.clone()), + ], + )?)?; + if exists.is_none() { + return Err(rejection( + py, + HistoryErrorKind::InvalidCursor, + UNKNOWN_RUN_CURSOR, + )); + } + Ok(Some(RunCursorKey { + created_at, + created_at_text: decoded.created_at_text, + run_id: decoded.run_id, + })) + } +} + +#[pymethods] +impl SqliteExecutionHistory { + #[new] + #[pyo3(signature = (*_args, **_kwargs))] + fn new(_args: &Bound<'_, PyTuple>, _kwargs: Option<&Bound<'_, PyDict>>) -> Self { + Self { + path: OnceLock::new(), + lock: OnceLock::new(), + closed: AtomicBool::new(true), + project_calls: AtomicI64::new(0), + } + } + + #[pyo3(signature = (project_dir = None, *, path = None, busy_timeout_ms = None))] + fn __init__( + slf: &Bound<'_, Self>, + project_dir: Option<&Bound<'_, PyAny>>, + path: Option<&Bound<'_, PyAny>>, + busy_timeout_ms: Option<&Bound<'_, PyAny>>, + ) -> PyResult<()> { + let py: Python<'_> = slf.py(); + let this: &Self = slf.get(); + let fspath = |value: Option<&Bound<'_, PyAny>>| -> PyResult> { + value + .filter(|value| !value.is_none()) + .map(|value| imported(py, "os", "fspath")?.call1((value,))?.extract()) + .transpose() + }; + let project_dir: Option = fspath(project_dir)?; + let path: Option = fspath(path)?; + let working_directory: String = if project_dir.is_none() && path.is_none() { + std::env::current_dir()?.to_string_lossy().into_owned() + } else { + String::new() + }; + let location = SqliteLocation { + project_dir, + path, + working_directory, + busy_timeout_ms: busy_timeout_ms.map_or(Some(SQLITE_BUSY_TIMEOUT_MS), exact_int), + }; + let database: String = checked(py, sqlite_database_path(&location))?; + let busy_timeout_ms: i64 = location.busy_timeout_ms.unwrap_or(SQLITE_BUSY_TIMEOUT_MS); + let _ = this.path.set(database.clone()); + let lock: Bound<'_, PyAny> = imported(py, THREADING_MODULE, "RLock")?.call0()?; + let _ = this.lock.set(lock.unbind()); + this.closed.store(false, Ordering::SeqCst); + let mut connection: Option> = None; + let opened: PyResult<()> = (|| { + prepare_history_directory(&database)?; + let sqlite = py.import(SQLITE_MODULE)?; + let opened: Bound<'_, PyAny> = sqlite.getattr("connect")?.call( + (database.as_str(),), + Some( + &crate::bindings::_helpers::compile_driver::python_calls::keywords( + py, + &[ + ( + "timeout", + (busy_timeout_ms as f64 / MILLIS_PER_SECOND) + .into_bound_py_any(py)?, + ), + ("isolation_level", py.None().into_bound(py)), + ("check_same_thread", false.into_bound_py_any(py)?), + ], + )?, + ), + )?; + connection = Some(opened.clone()); + slf.setattr("_connection", &opened)?; + this.with_lock(py, || { + opened.setattr("row_factory", sqlite.getattr("Row")?)?; + execute_plain(&opened, &format!("PRAGMA busy_timeout = {busy_timeout_ms}"))?; + execute_plain(&opened, SQLITE_FOREIGN_KEYS)?; + if database != MEMORY_PATH { + execute_plain(&opened, SQLITE_JOURNAL_MODE)?.call_method0("fetchone")?; + } + slf.call_method0("_migrate")?; + Ok(()) + })?; + restrict_history_file(&database)?; + slf.call_method0("reconcile")?; + Ok(()) + })(); + let Err(error) = opened else { + return Ok(()); + }; + let is_storage: bool = history_error_of_kind(py, &error, STORAGE_ERROR)?; + let is_open_failure: bool = + error.is_instance_of::(py) || sqlite_failure(py, &error)?; + if !is_storage && !is_open_failure { + return Err(error); + } + if let Some(connection) = connection { + let _ = connection.call_method0("close"); + } + this.closed.store(true, Ordering::SeqCst); + if is_storage { + return Err(error); + } + Err(storage_error_from( + py, + &format!("cannot open SQLite history at {database}"), + error, + )) + } + + /// The database path, or `None` for an in-memory database. + #[getter] + fn path<'py>(&self, py: Python<'py>) -> PyResult> { + let path: &str = self.database_path(); + if path == MEMORY_PATH { + return Ok(py.None().into_bound(py)); + } + imported(py, PATHLIB_MODULE, "Path")?.call1((path,)) + } + + /// The number of incremental projection calls. + #[getter] + fn project_calls(&self) -> i64 { + self.project_calls.load(Ordering::SeqCst) + } + + fn append_event<'py>( + slf: &Bound<'py, Self>, + event: &Bound<'py, PyAny>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + let this: &Self = slf.get(); + this.with_lock(py, || { + this.ensure_open(py)?; + slf.call_method1("append_events", (tuple_of(py, vec![event.clone()])?,))? + .get_item(0) + }) + } + + fn append_events<'py>( + slf: &Bound<'py, Self>, + events: &Bound<'py, PyAny>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + let this: &Self = slf.get(); + this.with_lock(py, || { + this.ensure_open(py)?; + let pending: Bound<'py, PyTuple> = + PyTuple::new(py, events.try_iter()?.collect::>>()?)?; + let appended = this.transaction(slf, |connection| { + Self::append_events_in(connection, pending.as_any()) + }); + let stored = Self::without_conflicts(py, appended, "SQLite event append failed")?; + tuple_of(py, stored) + }) + } + + #[pyo3(signature = (*, event_filter, after_cursor = None, limit = None))] + #[pyo3(name = "get_events")] + fn stored_events_page<'py>( + slf: &Bound<'py, Self>, + event_filter: &Bound<'py, PyAny>, + after_cursor: Option<&Bound<'py, PyAny>>, + limit: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + let this: &Self = slf.get(); + this.with_lock(py, || { + this.ensure_open(py)?; + let limit: Option = limit.map_or(Some(DEFAULT_PAGE_LIMIT), exact_int); + checked(py, validate_page_limit(limit))?; + let limit: i64 = limit.unwrap_or(DEFAULT_PAGE_LIMIT); + let connection: Bound<'py, PyAny> = slf.getattr("_connection")?; + let after: i64 = Self::event_position(&connection, after_cursor)?; + let spec = event_filter_spec(event_filter)?; + let query = event_page_query(Dialect::Sqlite, after, &spec, limit); + let rows: Vec> = wrapped( + py, + execute_query(&connection, &query) + .and_then(|cursor| cursor.try_iter()?.collect::>>()), + "SQLite event read failed", + )?; + let has_more: bool = rows.len() as i64 > limit; + let mut records: Vec> = Vec::new(); + if !rows.is_empty() { + let namespace: String = Self::namespace(&connection)?; + for row in rows.iter().take(limit as usize) { + records.push(Self::stored_event(row, &namespace)?); + } + } + let next_cursor: Bound<'py, PyAny> = match records.last() { + Some(record) => record.getattr("cursor")?, + None => py.None().into_bound(py), + }; + crate::bindings::_helpers::compile_driver::python_calls::call_with( + &imported( + py, + crate::bindings::_helpers::execution_history::conversion::MODELS_MODULE, + "EventPage", + )?, + &[ + ("records", tuple_of(py, records)?.into_any()), + ("next_cursor", next_cursor), + ("has_more", has_more.into_bound_py_any(py)?), + ], + ) + }) + } + + #[pyo3(name = "get_run")] + fn run_record<'py>( + slf: &Bound<'py, Self>, + run_id: &Bound<'py, PyAny>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + let this: &Self = slf.get(); + this.with_lock(py, || { + this.ensure_open(py)?; + let connection: Bound<'py, PyAny> = slf.getattr("_connection")?; + let row = + fetch_one(&connection.call_method1("execute", (SQLITE_RUN_BY_ID, (run_id,)))?)?; + match row { + Some(row) => crate::bindings::_helpers::execution_history::conversion::run_record( + py, + &stored_run(&row)?, + ), + None => Ok(py.None().into_bound(py)), + } + }) + } + + #[pyo3(signature = (*, run_filter, after_cursor = None, limit = None))] + #[pyo3(name = "get_runs")] + fn run_page<'py>( + slf: &Bound<'py, Self>, + run_filter: &Bound<'py, PyAny>, + after_cursor: Option<&Bound<'py, PyAny>>, + limit: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + let this: &Self = slf.get(); + this.with_lock(py, || { + this.ensure_open(py)?; + let limit: Option = limit.map_or(Some(DEFAULT_PAGE_LIMIT), exact_int); + checked(py, validate_page_limit(limit))?; + let limit: i64 = limit.unwrap_or(DEFAULT_PAGE_LIMIT); + let connection: Bound<'py, PyAny> = slf.getattr("_connection")?; + let none = py.None().into_bound(py); + let after = Self::run_cursor_key(&connection, after_cursor.unwrap_or(&none))?; + let spec = run_filter_spec(run_filter)?; + let query = run_page_query(Dialect::Sqlite, after.as_ref(), &spec, limit); + let rows: Vec> = execute_query(&connection, &query)? + .try_iter()? + .collect::>()?; + let runs: Vec = rows + .iter() + .take(limit as usize) + .map(stored_run) + .collect::>()?; + let next_cursor: Bound<'py, PyAny> = match runs.last() { + Some(run) => { + let namespace: String = Self::namespace(&connection)?; + checked(py, run_cursor(Dialect::Sqlite, &namespace, run))? + .into_bound_py_any(py)? + } + None => py.None().into_bound(py), + }; + crate::bindings::_helpers::compile_driver::python_calls::call_with( + &imported( + py, + crate::bindings::_helpers::execution_history::conversion::MODELS_MODULE, + "RunPage", + )?, + &[ + ("records", run_records(py, &runs)?.into_any()), + ("next_cursor", next_cursor), + ( + "has_more", + (rows.len() as i64 > limit).into_bound_py_any(py)?, + ), + ], + ) + }) + } + + fn project<'py>( + slf: &Bound<'py, Self>, + stored_events: &Bound<'py, PyAny>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + let this: &Self = slf.get(); + this.with_lock(py, || { + this.ensure_open(py)?; + this.project_calls.fetch_add(1, Ordering::SeqCst); + let connection: Bound<'py, PyAny> = slf.getattr("_connection")?; + let current: Vec = Self::read_all_runs(&connection)?; + let facts = projection_facts(stored_events)?; + let projected: Vec = checked(py, project_runs(facts, current))?; + let records: Bound<'py, PyTuple> = run_records(py, &projected)?; + slf.call_method1("_publish_projection", (&records,))?; + Ok(records) + }) + } + + fn rebuild_from_events<'py>( + slf: &Bound<'py, Self>, + stored_events: &Bound<'py, PyAny>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + let this: &Self = slf.get(); + this.with_lock(py, || { + this.ensure_open(py)?; + let facts = projection_facts(stored_events)?; + let projected: Vec = checked(py, project_runs(facts, Vec::new()))?; + let records: Bound<'py, PyTuple> = run_records(py, &projected)?; + slf.call_method1("_publish_projection", (&records,))?; + Ok(records) + }) + } + + /// Atomically append immutable facts and publish their run projection. + fn append_and_project<'py>( + slf: &Bound<'py, Self>, + events: &Bound<'py, PyAny>, + ) -> PyResult> { + let py: Python<'py> = slf.py(); + let this: &Self = slf.get(); + this.with_lock(py, || { + this.ensure_open(py)?; + let pending: Bound<'py, PyTuple> = + PyTuple::new(py, events.try_iter()?.collect::>>()?)?; + let appended = this.transaction(slf, |connection| { + let stored: Vec> = + Self::append_events_in(connection, pending.as_any())?; + let stored_tuple: Bound<'py, PyTuple> = tuple_of(py, stored.clone())?; + let facts = projection_facts(stored_tuple.as_any())?; + let projected = checked(py, project_runs(facts, Self::read_all_runs(connection)?))?; + Self::replace_projection(connection, &projected)?; + Ok(stored) + }); + let stored = Self::without_conflicts( + py, + appended, + "SQLite atomic append and projection failed", + )?; + tuple_of(py, stored) + }) + } + + /// Rebuild the disposable run projection from all durable event facts. + fn reconcile<'py>(slf: &Bound<'py, Self>) -> PyResult> { + let py: Python<'py> = slf.py(); + let this: &Self = slf.get(); + this.with_lock(py, || { + this.ensure_open(py)?; + let reconciled = this.transaction(slf, |connection| { + let records: Bound<'py, PyAny> = slf.call_method0("_read_all_events")?; + let facts = projection_facts(&records)?; + let projected: Vec = checked(py, project_runs(facts, Vec::new()))?; + Self::replace_projection(connection, &projected)?; + run_records(py, &projected) + }); + wrapped(py, reconciled, "SQLite reconciliation failed") + }) + } + + /// Verify connectivity, integrity, and the supported schema revision. + fn check_health(slf: &Bound<'_, Self>) -> PyResult { + let py: Python<'_> = slf.py(); + let this: &Self = slf.get(); + this.with_lock(py, || { + this.ensure_open(py)?; + let connection: Bound<'_, PyAny> = slf.getattr("_connection")?; + let Some(result) = fetch_one(&execute_plain(&connection, SQLITE_QUICK_CHECK)?)? else { + return Ok(false); + }; + if text_of(&result.get_item(0)?)?.as_deref() != Some(SQLITE_HEALTHY_CHECK_RESULT) { + return Ok(false); + } + let version: i64 = slf.call_method0("get_schema_version")?.extract()?; + Ok(version == SCHEMA_VERSION) + }) + } + + #[pyo3(name = "get_schema_version")] + fn schema_version(slf: &Bound<'_, Self>) -> PyResult { + let py: Python<'_> = slf.py(); + let this: &Self = slf.get(); + this.with_lock(py, || { + this.ensure_open(py)?; + Self::metadata_version(&slf.getattr("_connection")?) + }) + } + + #[pyo3(signature = (*, target_version = None))] + fn upgrade_schema(slf: &Bound<'_, Self>, target_version: Option) -> PyResult { + let py: Python<'_> = slf.py(); + let this: &Self = slf.get(); + this.with_lock(py, || { + this.ensure_open(py)?; + let target: i64 = target_version.unwrap_or(SCHEMA_VERSION); + checked(py, validate_schema_target(Dialect::Sqlite, target))?; + slf.call_method0("_migrate")?; + Ok(target) + }) + } + + fn close(slf: &Bound<'_, Self>) -> PyResult<()> { + let py: Python<'_> = slf.py(); + let this: &Self = slf.get(); + if this.lock.get().is_none() { + return Ok(()); + } + this.with_lock(py, || { + if this.closed.load(Ordering::SeqCst) { + return Ok(()); + } + let closed = slf + .getattr("_connection") + .and_then(|connection| connection.call_method0("close")); + this.closed.store(true, Ordering::SeqCst); + closed.map(|_| ()) + }) + } + + fn dispose(slf: &Bound<'_, Self>) -> PyResult<()> { + slf.call_method0("close")?; + Ok(()) + } + + fn __enter__<'py>(slf: &Bound<'py, Self>) -> PyResult> { + let py: Python<'py> = slf.py(); + let this: &Self = slf.get(); + this.with_lock(py, || this.ensure_open(py))?; + Ok(slf.clone()) + } + + #[pyo3(signature = (*_exc_info))] + fn __exit__(slf: &Bound<'_, Self>, _exc_info: &Bound<'_, PyTuple>) -> PyResult> { + slf.call_method0("close")?; + Ok(None) + } + + fn _migrate(slf: &Bound<'_, Self>) -> PyResult<()> { + slf.get().migrate_in(slf) + } + + fn _read_all_events<'py>(slf: &Bound<'py, Self>) -> PyResult> { + let py: Python<'py> = slf.py(); + let connection: Bound<'py, PyAny> = slf.getattr("_connection")?; + let rows: Vec> = execute_plain(&connection, SQLITE_ALL_EVENTS)? + .try_iter()? + .collect::>()?; + let mut records: Vec> = Vec::with_capacity(rows.len()); + if !rows.is_empty() { + let namespace: String = Self::namespace(&connection)?; + for row in &rows { + records.push(Self::stored_event(row, &namespace)?); + } + } + tuple_of(py, records) + } + + fn _publish_projection(slf: &Bound<'_, Self>, projected: &Bound<'_, PyAny>) -> PyResult<()> { + let py: Python<'_> = slf.py(); + let runs: Vec = run_rows(projected)?; + let published = slf.get().transaction(slf, |connection| { + Self::replace_projection(connection, &runs) + }); + wrapped(py, published, "SQLite run projection publication failed") + } +} + +impl SqliteExecutionHistory { + /// `result` with integrity conflicts kept and `sqlite3.Error` wrapped as `message`. + fn without_conflicts(py: Python<'_>, result: PyResult, message: &str) -> PyResult { + match result { + Err(error) if history_error_of_kind(py, &error, INTEGRITY_CONFLICT)? => Err(error), + result => wrapped(py, result, message), + } + } + + fn event_position( + connection: &Bound<'_, PyAny>, + cursor: Option<&Bound<'_, PyAny>>, + ) -> PyResult { + let py: Python<'_> = connection.py(); + let Some(cursor) = cursor.filter(|cursor| !cursor.is_none()) else { + return Ok(0); + }; + let text: String = text_of(cursor)?.unwrap_or_default(); + let namespace: String = Self::namespace(connection)?; + let position: i64 = checked( + py, + event_cursor_position(Dialect::Sqlite, Some(&text), &namespace), + )?; + let exists = fetch_one(&execute( + connection, + SQLITE_EVENT_EXISTS, + &[SqlParam::Integer(position)], + )?)?; + if exists.is_none() { + return Err(rejection( + py, + HistoryErrorKind::InvalidCursor, + UNKNOWN_EVENT_CURSOR, + )); + } + Ok(position) + } +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/execution_history/statements.rs b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/statements.rs new file mode 100644 index 0000000000..54e2c571f4 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/statements.rs @@ -0,0 +1,70 @@ +//! Running execution history SQL through a Python DB-API connection or cursor. + +use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyResult, Python}; +use pyo3::types::PyTuple; +use sqlbuild_state::execution_history::models::{Query, SqlParam}; + +use crate::bindings::_helpers::execution_history::conversion::python_params; + +/// `target.execute(sql, params)`. +pub(crate) fn execute<'py>( + target: &Bound<'py, PyAny>, + sql: &str, + params: &[SqlParam], +) -> PyResult> { + let py: Python<'py> = target.py(); + target.call_method1("execute", (sql, python_params(py, params)?)) +} + +/// `target.execute(sql)` without parameters. +pub(crate) fn execute_plain<'py>( + target: &Bound<'py, PyAny>, + sql: &str, +) -> PyResult> { + target.call_method1("execute", (sql,)) +} + +/// `target.execute(query.sql, query.params)`. +pub(crate) fn execute_query<'py>( + target: &Bound<'py, PyAny>, + query: &Query, +) -> PyResult> { + execute(target, &query.sql, &query.params) +} + +/// The first row `cursor.fetchone()` returns, or `None`. +pub(crate) fn fetch_one<'py>(cursor: &Bound<'py, PyAny>) -> PyResult>> { + let row: Bound<'py, PyAny> = cursor.call_method0("fetchone")?; + Ok((!row.is_none()).then_some(row)) +} + +/// Every row `cursor.fetchall()` returns. +pub(crate) fn fetch_all<'py>(cursor: &Bound<'py, PyAny>) -> PyResult>> { + cursor.call_method0("fetchall")?.try_iter()?.collect() +} + +/// Run `body` while holding the Python `lock`, releasing it on every exit. +pub(crate) fn locked( + lock: &Bound<'_, PyAny>, + body: impl FnOnce() -> PyResult, +) -> PyResult { + lock.call_method0("acquire")?; + let result: PyResult = body(); + let released = lock.call_method0("release"); + let value: T = result?; + released?; + Ok(value) +} + +/// `row[name]`. +pub(crate) fn column<'py>(row: &Bound<'py, PyAny>, name: &str) -> PyResult> { + row.get_item(name) +} + +/// The tuple of `items`. +pub(crate) fn tuple_of<'py>( + py: Python<'py>, + items: Vec>, +) -> PyResult> { + PyTuple::new(py, items) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/mod.rs index 5d9d284f4b..3277ea6a3a 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/mod.rs @@ -32,6 +32,7 @@ pub(crate) mod debug_command; pub(crate) mod diagnostics; pub(crate) mod diff_report; pub(crate) mod discovery; +pub(crate) mod execution_history; pub(crate) mod fingerprints; pub(crate) mod format_command; pub(crate) mod freshness; diff --git a/crates/sqlbuild-state/src/execution_history/_helpers/base64url.rs b/crates/sqlbuild-state/src/execution_history/_helpers/base64url.rs new file mode 100644 index 0000000000..7dcc81618c --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/_helpers/base64url.rs @@ -0,0 +1,70 @@ +//! URL-safe base64 as Python's `base64.urlsafe_b64encode` and `urlsafe_b64decode` apply it. + +const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_"; +const PAD: u8 = b'='; +const SEXTET_BITS: u32 = 6; +const BYTE_BITS: u32 = 8; +const QUAD_LENGTH: usize = 4; +const MAX_PADDING: usize = 2; +const TRIPLE_LENGTH: usize = 3; +const LOW_SEXTET: u32 = 0x3f; +const LOW_BYTE: u32 = 0xff; + +/// `base64.urlsafe_b64encode(value).decode()`. +pub(crate) fn encode(value: &[u8]) -> String { + let mut encoded = String::with_capacity(value.len().div_ceil(TRIPLE_LENGTH) * QUAD_LENGTH); + for chunk in value.chunks(TRIPLE_LENGTH) { + let bits: u32 = chunk + .iter() + .chain(std::iter::repeat_n(&0, TRIPLE_LENGTH - chunk.len())) + .fold(0, |bits, byte| (bits << BYTE_BITS) | u32::from(*byte)); + for position in 0..QUAD_LENGTH { + if position > chunk.len() { + encoded.push(char::from(PAD)); + continue; + } + let shift: u32 = SEXTET_BITS * (3 - position as u32); + encoded.push(char::from( + ALPHABET[((bits >> shift) & LOW_SEXTET) as usize], + )); + } + } + encoded +} + +/// Decoded bytes of padded URL-safe or standard base64, or `None` when malformed. +pub(crate) fn decode(text: &str) -> Option> { + let bytes: &[u8] = text.as_bytes(); + if !bytes.len().is_multiple_of(QUAD_LENGTH) { + return None; + } + let mut decoded: Vec = Vec::with_capacity(bytes.len() / QUAD_LENGTH * TRIPLE_LENGTH); + let quad_count: usize = bytes.len() / QUAD_LENGTH; + for (index, quad) in bytes.chunks(QUAD_LENGTH).enumerate() { + let padding: usize = quad.iter().rev().take_while(|byte| **byte == PAD).count(); + if padding > MAX_PADDING || (padding > 0 && index + 1 != quad_count) { + return None; + } + let mut bits: u32 = 0; + for byte in &quad[..QUAD_LENGTH - padding] { + bits = (bits << SEXTET_BITS) | u32::from(sextet(*byte)?); + } + bits <<= SEXTET_BITS * padding as u32; + for position in 0..(TRIPLE_LENGTH - padding) { + let shift: u32 = BYTE_BITS * (2 - position as u32); + decoded.push(((bits >> shift) & LOW_BYTE) as u8); + } + } + Some(decoded) +} + +fn sextet(character: u8) -> Option { + match character { + b'A'..=b'Z' => Some(character - b'A'), + b'a'..=b'z' => Some(character - b'a' + 26), + b'0'..=b'9' => Some(character - b'0' + 52), + b'-' | b'+' => Some(62), + b'_' | b'/' => Some(63), + _ => None, + } +} diff --git a/crates/sqlbuild-state/src/execution_history/_helpers/dialects.rs b/crates/sqlbuild-state/src/execution_history/_helpers/dialects.rs new file mode 100644 index 0000000000..67e055005f --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/_helpers/dialects.rs @@ -0,0 +1,68 @@ +//! Per-dialect placeholders, tables, cursor prefixes and parameter values. + +use crate::execution_history::_helpers::timestamps::stored_text; +use crate::execution_history::constants::{ + POSTGRES_EVENT_CURSOR_PREFIX, POSTGRES_EVENT_TABLE, POSTGRES_PLACEHOLDER, + POSTGRES_RUN_CURSOR_PREFIX, POSTGRES_RUN_TABLE, SQLITE_EVENT_CURSOR_PREFIX, SQLITE_EVENT_TABLE, + SQLITE_PLACEHOLDER, SQLITE_RUN_CURSOR_PREFIX, SQLITE_RUN_TABLE, +}; +use crate::execution_history::models::{Dialect, SqlParam}; +use crate::execution_history::types::UtcTime; + +/// The parameter placeholder. +pub(crate) fn placeholder(dialect: Dialect) -> &'static str { + match dialect { + Dialect::Sqlite => SQLITE_PLACEHOLDER, + Dialect::Postgres => POSTGRES_PLACEHOLDER, + } +} + +/// The lifecycle event log table. +pub(crate) fn event_table(dialect: Dialect) -> &'static str { + match dialect { + Dialect::Sqlite => SQLITE_EVENT_TABLE, + Dialect::Postgres => POSTGRES_EVENT_TABLE, + } +} + +/// The run projection table. +pub(crate) fn run_table(dialect: Dialect) -> &'static str { + match dialect { + Dialect::Sqlite => SQLITE_RUN_TABLE, + Dialect::Postgres => POSTGRES_RUN_TABLE, + } +} + +/// The event cursor prefix. +pub(crate) fn event_cursor_prefix(dialect: Dialect) -> &'static str { + match dialect { + Dialect::Sqlite => SQLITE_EVENT_CURSOR_PREFIX, + Dialect::Postgres => POSTGRES_EVENT_CURSOR_PREFIX, + } +} + +/// The run cursor prefix. +pub(crate) fn run_cursor_prefix(dialect: Dialect) -> &'static str { + match dialect { + Dialect::Sqlite => SQLITE_RUN_CURSOR_PREFIX, + Dialect::Postgres => POSTGRES_RUN_CURSOR_PREFIX, + } +} + +/// A timestamp parameter: SQLite stores ISO text, PostgreSQL binds a `datetime`. +pub(crate) fn timestamp_param(dialect: Dialect, value: &UtcTime) -> SqlParam { + match dialect { + Dialect::Sqlite => SqlParam::Text(stored_text(value)), + Dialect::Postgres => SqlParam::Timestamp(*value), + } +} + +/// An optional timestamp parameter. +pub(crate) fn optional_timestamp_param(dialect: Dialect, value: Option<&UtcTime>) -> SqlParam { + value.map_or(SqlParam::Null, |value| timestamp_param(dialect, value)) +} + +/// An optional text parameter. +pub(crate) fn optional_text_param(value: Option<&String>) -> SqlParam { + value.map_or(SqlParam::Null, |value| SqlParam::Text(value.clone())) +} diff --git a/crates/sqlbuild-state/src/execution_history/_helpers/filters.rs b/crates/sqlbuild-state/src/execution_history/_helpers/filters.rs new file mode 100644 index 0000000000..8de3e7cf02 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/_helpers/filters.rs @@ -0,0 +1,40 @@ +//! Filter field checks shared by event and run filters. + +use sqlbuild_core::text::main::python_strip::python_strip; + +use crate::execution_history::errors::{HistoryError, HistoryErrorKind}; +use crate::execution_history::models::FilterTime; + +/// An `InvalidFilterError`. +pub(crate) fn filter_error(message: &str) -> HistoryError { + HistoryError::new(HistoryErrorKind::InvalidFilter, message) +} + +/// An optional correlation value must not be blank. +pub(crate) fn filled_text(value: Option<&str>, field_name: &str) -> Result<(), HistoryError> { + match value { + Some(value) if python_strip(value).is_empty() => Err(filter_error(&format!( + "{field_name} must be non-empty when provided" + ))), + _ => Ok(()), + } +} + +/// An optional bound must use `UTC`. +pub(crate) fn utc_bound(value: Option<&FilterTime>, field_name: &str) -> Result<(), HistoryError> { + match value { + Some(value) if !value.is_utc => Err(filter_error(&format!("{field_name} must use UTC"))), + _ => Ok(()), + } +} + +/// Both bounds, when present, must not be reversed. +pub(crate) fn bounds_ordered( + (start, end): (Option<&FilterTime>, Option<&FilterTime>), + message: &str, +) -> Result<(), HistoryError> { + match (start, end) { + (Some(start), Some(end)) if start.instant > end.instant => Err(filter_error(message)), + _ => Ok(()), + } +} diff --git a/crates/sqlbuild-state/src/execution_history/_helpers/mod.rs b/crates/sqlbuild-state/src/execution_history/_helpers/mod.rs new file mode 100644 index 0000000000..17b8c848bc --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/_helpers/mod.rs @@ -0,0 +1,5 @@ +pub(crate) mod base64url; +pub(crate) mod dialects; +pub(crate) mod filters; +pub(crate) mod schema; +pub(crate) mod timestamps; diff --git a/crates/sqlbuild-state/src/execution_history/_helpers/schema.rs b/crates/sqlbuild-state/src/execution_history/_helpers/schema.rs new file mode 100644 index 0000000000..423670998a --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/_helpers/schema.rs @@ -0,0 +1,16 @@ +//! Schema version rejections. + +use crate::execution_history::errors::{HistoryError, HistoryErrorKind}; +use crate::execution_history::models::Dialect; + +/// The rejection of an unsupported requested or stored schema version. +pub(crate) fn unsupported_version(dialect: Dialect, version: i64) -> HistoryError { + let backend: &str = match dialect { + Dialect::Sqlite => "SQLite", + Dialect::Postgres => "PostgreSQL execution history", + }; + HistoryError::new( + HistoryErrorKind::UnsupportedSchemaVersion, + format!("unsupported {backend} schema version {version}"), + ) +} diff --git a/crates/sqlbuild-state/src/execution_history/_helpers/timestamps.rs b/crates/sqlbuild-state/src/execution_history/_helpers/timestamps.rs new file mode 100644 index 0000000000..e906ff83b7 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/_helpers/timestamps.rs @@ -0,0 +1,29 @@ +//! ISO text of UTC instants as Python's `datetime.isoformat` writes them. + +use crate::execution_history::constants::UTC_OFFSET_TEXT; +use crate::execution_history::types::UtcTime; + +/// `value.isoformat(timespec="microseconds")` of a UTC `datetime`. +pub(crate) fn stored_text(value: &UtcTime) -> String { + format!( + "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}.{:06}{UTC_OFFSET_TEXT}", + value.date.year, + value.date.month, + value.date.day, + value.hour, + value.minute, + value.second, + value.microsecond + ) +} + +/// `value.isoformat()` of a UTC `datetime`: microseconds only when non-zero. +pub(crate) fn default_text(value: &UtcTime) -> String { + if value.microsecond != 0 { + return stored_text(value); + } + format!( + "{:04}-{:02}-{:02}T{:02}:{:02}:{:02}{UTC_OFFSET_TEXT}", + value.date.year, value.date.month, value.date.day, value.hour, value.minute, value.second + ) +} diff --git a/crates/sqlbuild-state/src/execution_history/constants.rs b/crates/sqlbuild-state/src/execution_history/constants.rs new file mode 100644 index 0000000000..d358f18c0e --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/constants.rs @@ -0,0 +1,136 @@ +//! Execution history limits, lifecycle facts and the SQL each backend stores them with. + +pub const DEFAULT_PAGE_LIMIT: i64 = 100; +pub const MAX_PAGE_LIMIT: i64 = 1000; +pub const SCHEMA_VERSION: i64 = 1; +pub const RUN_STARTED_EVENT_TYPE: &str = "run_started"; +pub const RUN_COMPLETED_EVENT_TYPE: &str = "run_completed"; +pub const RUN_FAILED_EVENT_TYPE: &str = "run_failed"; +pub const RUN_CURSOR_PART_COUNT: usize = 3; +pub const SQLITE_EVENT_CURSOR_PREFIX: &str = "sqlite-event:"; +pub const SQLITE_RUN_CURSOR_PREFIX: &str = "sqlite-run:"; +pub const POSTGRES_EVENT_CURSOR_PREFIX: &str = "postgres-event:"; +pub const POSTGRES_RUN_CURSOR_PREFIX: &str = "postgres-run:"; +pub const SQLITE_BUSY_TIMEOUT_MS: i64 = 5_000; +pub const POSTGRES_CONNECT_TIMEOUT_SECONDS: i64 = 10; +pub const POSTGRES_STATEMENT_TIMEOUT_MS: i64 = 30_000; +pub const POSTGRES_TRANSACTION_RETRIES: i64 = 3; +pub const POSTGRES_MAX_CONNECT_TIMEOUT_SECONDS: i64 = 300; +pub const POSTGRES_MAX_STATEMENT_TIMEOUT_MS: i64 = 3_600_000; +pub const POSTGRES_MAX_TRANSACTION_RETRIES: i64 = 10; +pub const POSTGRES_MIGRATION_LOCK_ID: i64 = 7_218_786_101; +pub const POSTGRES_HISTORY_LOCK_ID: i64 = 7_218_786_102; +pub const POSTGRES_RETRYABLE_SQLSTATES: [&str; 2] = ["40001", "40P01"]; +/// Text envelope fields stored next to each event, in `EventEnvelope::text_fields` order. +pub const ENVELOPE_TEXT_FIELDS: [&str; 9] = [ + "producer", + "producer_version", + "event_type", + "invocation_id", + "run_id", + "resource_id", + "resource_attempt_id", + "operation_id", + "statement_id", +]; +pub const SQLITE_SCHEMA: [&str; 8] = [ + "CREATE TABLE IF NOT EXISTS execution_history_metadata (\n key TEXT PRIMARY KEY,\n value TEXT NOT NULL\n)", + "CREATE TABLE IF NOT EXISTS lifecycle_event_log (\n storage_id INTEGER PRIMARY KEY AUTOINCREMENT,\n event_id TEXT NOT NULL UNIQUE,\n schema_version INTEGER NOT NULL,\n producer TEXT,\n event_type TEXT,\n occurred_at TEXT,\n received_at TEXT NOT NULL,\n invocation_id TEXT,\n run_id TEXT,\n resource_id TEXT,\n resource_attempt_id TEXT,\n operation_id TEXT,\n statement_id TEXT,\n payload_json TEXT NOT NULL,\n content_digest TEXT NOT NULL\n)", + "CREATE INDEX IF NOT EXISTS lifecycle_event_log_run_storage\n ON lifecycle_event_log (run_id, storage_id)", + "CREATE INDEX IF NOT EXISTS lifecycle_event_log_invocation_storage\n ON lifecycle_event_log (invocation_id, storage_id)", + "CREATE INDEX IF NOT EXISTS lifecycle_event_log_type_storage\n ON lifecycle_event_log (event_type, storage_id)", + "CREATE TABLE IF NOT EXISTS run_projection (\n run_id TEXT PRIMARY KEY,\n invocation_id TEXT NOT NULL,\n created_at TEXT NOT NULL,\n status TEXT NOT NULL,\n is_complete INTEGER NOT NULL,\n last_event_cursor TEXT NOT NULL,\n last_storage_order INTEGER NOT NULL,\n command TEXT,\n target TEXT,\n environment TEXT,\n started_at TEXT,\n ended_at TEXT,\n projection_schema_version INTEGER NOT NULL\n)", + "CREATE INDEX IF NOT EXISTS run_projection_created ON run_projection (created_at, run_id)", + "CREATE INDEX IF NOT EXISTS run_projection_invocation_created\n ON run_projection (invocation_id, created_at, run_id)", +]; +pub const POSTGRES_SCHEMA: [&str; 8] = [ + "CREATE TABLE IF NOT EXISTS sqlbuild_storage_migrations (\n singleton BOOLEAN PRIMARY KEY DEFAULT TRUE CHECK (singleton),\n schema_version INTEGER NOT NULL,\n storage_namespace TEXT NOT NULL,\n migrated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP\n )", + "CREATE TABLE IF NOT EXISTS sqlbuild_lifecycle_event_log (\n storage_id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY,\n event_id TEXT NOT NULL UNIQUE,\n schema_version INTEGER NOT NULL,\n producer TEXT,\n producer_version TEXT,\n event_type TEXT,\n occurred_at TIMESTAMPTZ,\n received_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,\n invocation_id TEXT,\n run_id TEXT,\n resource_id TEXT,\n resource_attempt_id TEXT,\n operation_id TEXT,\n statement_id TEXT,\n payload_json TEXT NOT NULL,\n content_digest TEXT NOT NULL\n )", + "CREATE INDEX IF NOT EXISTS sqlbuild_lifecycle_event_log_run_storage\n ON sqlbuild_lifecycle_event_log (run_id, storage_id)", + "CREATE INDEX IF NOT EXISTS sqlbuild_lifecycle_event_log_invocation_storage\n ON sqlbuild_lifecycle_event_log (invocation_id, storage_id)", + "CREATE INDEX IF NOT EXISTS sqlbuild_lifecycle_event_log_type_storage\n ON sqlbuild_lifecycle_event_log (event_type, storage_id)", + "CREATE TABLE IF NOT EXISTS sqlbuild_run_projection (\n run_id TEXT PRIMARY KEY,\n invocation_id TEXT NOT NULL,\n created_at TIMESTAMPTZ NOT NULL,\n status TEXT NOT NULL,\n is_complete BOOLEAN NOT NULL,\n last_event_cursor TEXT NOT NULL,\n last_event_storage_id BIGINT NOT NULL,\n command TEXT,\n target TEXT,\n environment TEXT,\n started_at TIMESTAMPTZ,\n ended_at TIMESTAMPTZ,\n projection_schema_version INTEGER NOT NULL\n )", + "CREATE INDEX IF NOT EXISTS sqlbuild_run_projection_created\n ON sqlbuild_run_projection (created_at, run_id)", + "CREATE INDEX IF NOT EXISTS sqlbuild_run_projection_invocation_created\n ON sqlbuild_run_projection (invocation_id, created_at, run_id)", +]; +pub const SQLITE_EVENT_INSERT: &str = "INSERT INTO lifecycle_event_log (\n event_id, schema_version, producer, event_type, occurred_at, received_at,\n invocation_id, run_id, resource_id, resource_attempt_id, operation_id,\n statement_id, payload_json, content_digest\n ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"; +pub const POSTGRES_EVENT_INSERT: &str = "INSERT INTO sqlbuild_lifecycle_event_log (\n event_id, schema_version, producer, producer_version, event_type, occurred_at,\n invocation_id, run_id, resource_id, resource_attempt_id, operation_id,\n statement_id, payload_json, content_digest\n ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)\n ON CONFLICT (event_id) DO NOTHING RETURNING *"; +pub const SQLITE_RUN_INSERT: &str = "INSERT INTO run_projection (\n run_id, invocation_id, created_at, status, is_complete, last_event_cursor,\n last_storage_order, command, target, environment, started_at, ended_at,\n projection_schema_version\n ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"; +pub const POSTGRES_RUN_INSERT: &str = "INSERT INTO sqlbuild_run_projection (\n run_id, invocation_id, created_at, status, is_complete, last_event_cursor,\n last_event_storage_id, command, target, environment, started_at, ended_at,\n projection_schema_version\n ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)"; +pub const MEMORY_PATH: &str = ":memory:"; +pub const SQLITE_DIRECTORY_NAME: &str = ".sqlbuild"; +pub const SQLITE_FILE_NAME: &str = "history.sqlite3"; +pub const SQLITE_DIRECTORY_MODE: u32 = 0o700; +pub const SQLITE_FILE_MODE: u32 = 0o600; +pub const SQLITE_HEALTHY_CHECK_RESULT: &str = "ok"; +pub const SQLITE_METADATA_TABLE: &str = "execution_history_metadata"; +pub const SQLITE_NAMESPACE_KEY: &str = "storage_namespace"; +pub const SQLITE_SCHEMA_VERSION_KEY: &str = "schema_version"; +pub const SQLITE_FOREIGN_KEYS: &str = "PRAGMA foreign_keys = ON"; +pub const SQLITE_JOURNAL_MODE: &str = "PRAGMA journal_mode = WAL"; +pub const SQLITE_QUICK_CHECK: &str = "PRAGMA quick_check"; +pub const SQLITE_BEGIN: &str = "BEGIN IMMEDIATE"; +pub const SQLITE_COMMIT: &str = "COMMIT"; +pub const SQLITE_ROLLBACK: &str = "ROLLBACK"; +pub const SQLITE_INSERT_NAMESPACE: &str = + "INSERT OR IGNORE INTO execution_history_metadata(key, value) VALUES (?, ?)"; +pub const SQLITE_INSERT_SCHEMA_VERSION: &str = + "INSERT OR REPLACE INTO execution_history_metadata(key, value) VALUES (?, ?)"; +pub const SQLITE_METADATA_EXISTS: &str = + "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?"; +pub const SQLITE_METADATA_VALUE: &str = + "SELECT value FROM execution_history_metadata WHERE key = ?"; +pub const SQLITE_EVENT_BY_ID: &str = "SELECT * FROM lifecycle_event_log WHERE event_id = ?"; +pub const SQLITE_EVENT_BY_STORAGE_ID: &str = + "SELECT * FROM lifecycle_event_log WHERE storage_id = ?"; +pub const SQLITE_EVENT_EXISTS: &str = "SELECT 1 FROM lifecycle_event_log WHERE storage_id = ?"; +pub const SQLITE_ALL_EVENTS: &str = "SELECT * FROM lifecycle_event_log ORDER BY storage_id"; +pub const SQLITE_ALL_RUNS: &str = "SELECT * FROM run_projection ORDER BY created_at, run_id"; +pub const SQLITE_RUN_BY_ID: &str = "SELECT * FROM run_projection WHERE run_id = ?"; +pub const SQLITE_RUN_EXISTS: &str = + "SELECT 1 FROM run_projection WHERE created_at = ? AND run_id = ?"; +pub const SQLITE_DELETE_RUNS: &str = "DELETE FROM run_projection"; +pub const POSTGRES_SESSION_TIMEOUT: &str = "SELECT set_config('statement_timeout', %s, false)"; +pub const POSTGRES_TRANSACTION_TIMEOUT: &str = "SELECT set_config('statement_timeout', %s, true)"; +pub const POSTGRES_ADVISORY_LOCK: &str = "SELECT pg_advisory_xact_lock(%s)"; +pub const POSTGRES_MIGRATIONS_EXIST: &str = + "SELECT to_regclass('sqlbuild_storage_migrations') AS table_name"; +pub const POSTGRES_SCHEMA_VERSION: &str = + "SELECT schema_version FROM sqlbuild_storage_migrations WHERE singleton = TRUE"; +pub const POSTGRES_NAMESPACE: &str = + "SELECT storage_namespace FROM sqlbuild_storage_migrations WHERE singleton = TRUE"; +pub const POSTGRES_RECORD_MIGRATION: &str = "INSERT INTO sqlbuild_storage_migrations\n (singleton, schema_version, storage_namespace)\n VALUES (TRUE, %s, %s)\n ON CONFLICT (singleton) DO UPDATE SET\n schema_version = EXCLUDED.schema_version,\n migrated_at = CURRENT_TIMESTAMP"; +pub const POSTGRES_EVENT_BY_ID: &str = + "SELECT * FROM sqlbuild_lifecycle_event_log WHERE event_id = %s"; +pub const POSTGRES_EVENT_EXISTS: &str = + "SELECT 1 AS present FROM sqlbuild_lifecycle_event_log WHERE storage_id = %s"; +pub const POSTGRES_ALL_EVENTS: &str = + "SELECT * FROM sqlbuild_lifecycle_event_log ORDER BY storage_id"; +pub const POSTGRES_ALL_RUNS: &str = + "SELECT * FROM sqlbuild_run_projection ORDER BY created_at, run_id"; +pub const POSTGRES_RUN_BY_ID: &str = "SELECT * FROM sqlbuild_run_projection WHERE run_id = %s"; +pub const POSTGRES_RUN_EXISTS: &str = + "SELECT 1 AS present FROM sqlbuild_run_projection WHERE created_at = %s AND run_id = %s"; +pub const POSTGRES_DELETE_RUNS: &str = "DELETE FROM sqlbuild_run_projection"; +pub const POSTGRES_PING: &str = "SELECT 1"; +pub const SQLITE_PLACEHOLDER: &str = "?"; +pub const POSTGRES_PLACEHOLDER: &str = "%s"; +pub const SQLITE_EVENT_TABLE: &str = "lifecycle_event_log"; +pub const POSTGRES_EVENT_TABLE: &str = "sqlbuild_lifecycle_event_log"; +pub const SQLITE_RUN_TABLE: &str = "run_projection"; +pub const POSTGRES_RUN_TABLE: &str = "sqlbuild_run_projection"; +pub const CURSOR_SEPARATOR: char = ':'; +pub const MALFORMED_EVENT_CURSOR: &str = "event cursor is malformed or belongs to another backend"; +pub const FOREIGN_EVENT_CURSOR: &str = "event cursor is not valid for this storage"; +pub const UNKNOWN_EVENT_CURSOR: &str = "event cursor does not identify a durable event"; +pub const MALFORMED_RUN_CURSOR: &str = "run cursor is malformed or belongs to another backend"; +pub const RUN_CURSOR_TIMESTAMP: &str = "run cursor contains an invalid timestamp"; +pub const UNKNOWN_RUN_CURSOR: &str = "run cursor does not identify a run in this storage"; +pub const MISSING_EVENT_ID: &str = "canonical lifecycle event requires a non-empty string event_id"; +pub const UTC_OFFSET_TEXT: &str = "+00:00"; +pub const FIRST_POSITION: i64 = 1; +pub const NO_SCHEMA: i64 = 0; +pub const MIN_TIMEOUT: i64 = 1; +pub const MIN_RETRIES: i64 = 0; +pub const MIN_PAGE_LIMIT: i64 = 1; +pub const DUPLICATE_ORDER: &str = "storage_order must uniquely identify a durable event"; diff --git a/crates/sqlbuild-state/src/execution_history/errors.rs b/crates/sqlbuild-state/src/execution_history/errors.rs new file mode 100644 index 0000000000..55721b1b40 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/errors.rs @@ -0,0 +1,31 @@ +//! Execution history rejections, each naming the Python `ExecutionHistoryStorageError` subclass. + +/// Which `sqlbuild.runtime.execution_history.exceptions` class a rejection raises. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum HistoryErrorKind { + Storage, + IntegrityConflict, + InvalidEvent, + InvalidCursor, + InvalidFilter, + InvalidLimit, + UnsupportedSchemaVersion, + ProjectionConsistency, +} + +/// One rejection and its message. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct HistoryError { + pub kind: HistoryErrorKind, + pub message: String, +} + +impl HistoryError { + /// A rejection of `kind` with `message`. + pub fn new(kind: HistoryErrorKind, message: impl Into) -> Self { + Self { + kind, + message: message.into(), + } + } +} diff --git a/crates/sqlbuild-state/src/execution_history/main/canonical_event_id.rs b/crates/sqlbuild-state/src/execution_history/main/canonical_event_id.rs new file mode 100644 index 0000000000..6f4d7324a7 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/canonical_event_id.rs @@ -0,0 +1,15 @@ +//! `canonical_event_id`: the required stable id of a known or opaque lifecycle event. + +use crate::execution_history::constants::MISSING_EVENT_ID; +use crate::execution_history::errors::{HistoryError, HistoryErrorKind}; + +/// The event's `event_id` when it is a non-empty string. +pub fn canonical_event_id(event_id: Option<&str>) -> Result { + match event_id { + Some(event_id) if !event_id.is_empty() => Ok(event_id.to_owned()), + _ => Err(HistoryError::new( + HistoryErrorKind::InvalidEvent, + MISSING_EVENT_ID, + )), + } +} diff --git a/crates/sqlbuild-state/src/execution_history/main/decoded_run_cursor.rs b/crates/sqlbuild-state/src/execution_history/main/decoded_run_cursor.rs new file mode 100644 index 0000000000..a09bfce65d --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/decoded_run_cursor.rs @@ -0,0 +1,39 @@ +//! `decoded_run_cursor`: the `created_at` text and run id a run cursor names. + +use serde_json::Value; + +use crate::execution_history::_helpers::base64url::decode; +use crate::execution_history::_helpers::dialects::run_cursor_prefix; +use crate::execution_history::constants::{MALFORMED_RUN_CURSOR, RUN_CURSOR_PART_COUNT}; +use crate::execution_history::errors::{HistoryError, HistoryErrorKind}; +use crate::execution_history::models::{DecodedRunCursor, Dialect}; + +/// The position of a well-formed run cursor issued by the storage with `namespace`. +pub fn decoded_run_cursor( + dialect: Dialect, + cursor: &str, + namespace: &str, +) -> Result { + let malformed = || HistoryError::new(HistoryErrorKind::InvalidCursor, MALFORMED_RUN_CURSOR); + let encoded: &str = cursor + .strip_prefix(run_cursor_prefix(dialect)) + .ok_or_else(malformed)?; + let bytes: Vec = decode(encoded).ok_or_else(malformed)?; + let text: String = String::from_utf8(bytes).map_err(|_| malformed())?; + let decoded: Value = serde_json::from_str(&text).map_err(|_| malformed())?; + let Value::Array(parts) = decoded else { + return Err(malformed()); + }; + let [cursor_namespace, created_at, run_id]: [Value; RUN_CURSOR_PART_COUNT] = + parts.try_into().map_err(|_| malformed())?; + let (Value::String(created_at), Value::String(run_id)) = (created_at, run_id) else { + return Err(malformed()); + }; + if cursor_namespace.as_str() != Some(namespace) || run_id.is_empty() { + return Err(malformed()); + } + Ok(DecodedRunCursor { + created_at_text: created_at, + run_id, + }) +} diff --git a/crates/sqlbuild-state/src/execution_history/main/event_cursor.rs b/crates/sqlbuild-state/src/execution_history/main/event_cursor.rs new file mode 100644 index 0000000000..b370474b59 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/event_cursor.rs @@ -0,0 +1,9 @@ +//! `event_cursor`: the opaque cursor of one durable event. + +use crate::execution_history::_helpers::dialects::event_cursor_prefix; +use crate::execution_history::models::Dialect; + +/// `:`. +pub fn event_cursor(dialect: Dialect, namespace: &str, storage_id: i64) -> String { + format!("{}{namespace}:{storage_id}", event_cursor_prefix(dialect)) +} diff --git a/crates/sqlbuild-state/src/execution_history/main/event_cursor_position.rs b/crates/sqlbuild-state/src/execution_history/main/event_cursor_position.rs new file mode 100644 index 0000000000..b80e8409e4 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/event_cursor_position.rs @@ -0,0 +1,36 @@ +//! `event_cursor_position`: the storage position an event cursor names. + +use crate::execution_history::_helpers::dialects::event_cursor_prefix; +use crate::execution_history::constants::{ + CURSOR_SEPARATOR, FIRST_POSITION, FOREIGN_EVENT_CURSOR, MALFORMED_EVENT_CURSOR, +}; +use crate::execution_history::errors::{HistoryError, HistoryErrorKind}; +use crate::execution_history::models::Dialect; + +/// Position zero without a cursor, else the positive position of this storage's cursor. +pub fn event_cursor_position( + dialect: Dialect, + cursor: Option<&str>, + namespace: &str, +) -> Result { + let Some(cursor) = cursor else { + return Ok(0); + }; + let malformed = || HistoryError::new(HistoryErrorKind::InvalidCursor, MALFORMED_EVENT_CURSOR); + let mut parts = cursor.rsplitn(3, CURSOR_SEPARATOR); + let (Some(raw_position), Some(cursor_namespace), Some(prefix)) = + (parts.next(), parts.next(), parts.next()) + else { + return Err(malformed()); + }; + let position: i64 = raw_position.trim().parse().map_err(|_| malformed())?; + let expected_prefix: &str = event_cursor_prefix(dialect); + let prefix_matches: bool = expected_prefix.strip_suffix(CURSOR_SEPARATOR) == Some(prefix); + if !prefix_matches || cursor_namespace != namespace || position < FIRST_POSITION { + return Err(HistoryError::new( + HistoryErrorKind::InvalidCursor, + FOREIGN_EVENT_CURSOR, + )); + } + Ok(position) +} diff --git a/crates/sqlbuild-state/src/execution_history/main/event_insert_query.rs b/crates/sqlbuild-state/src/execution_history/main/event_insert_query.rs new file mode 100644 index 0000000000..9553849a53 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/event_insert_query.rs @@ -0,0 +1,45 @@ +//! `event_insert_query`: the insert of one prepared event with its envelope. + +use crate::execution_history::_helpers::dialects::{ + optional_text_param, optional_timestamp_param, timestamp_param, +}; +use crate::execution_history::constants::{POSTGRES_EVENT_INSERT, SQLITE_EVENT_INSERT}; +use crate::execution_history::models::{Dialect, PreparedEvent, Query, SqlParam}; +use crate::execution_history::types::UtcTime; + +/// SQLite stores `received_at` from the caller; PostgreSQL stores `producer_version` instead. +pub fn event_insert_query(dialect: Dialect, event: &PreparedEvent, received_at: &UtcTime) -> Query { + let envelope = &event.envelope; + let [producer, producer_version, event_type, correlations @ ..] = &envelope.text_fields; + let mut params: Vec = vec![ + SqlParam::Text(event.event_id.clone()), + SqlParam::Integer(envelope.schema_version), + optional_text_param(producer.as_ref()), + ]; + if dialect == Dialect::Postgres { + params.push(optional_text_param(producer_version.as_ref())); + } + params.push(optional_text_param(event_type.as_ref())); + params.push(optional_timestamp_param( + dialect, + envelope.occurred_at.as_ref(), + )); + if dialect == Dialect::Sqlite { + params.push(timestamp_param(dialect, received_at)); + } + params.extend( + correlations + .iter() + .map(|value| optional_text_param(value.as_ref())), + ); + params.push(SqlParam::Text(event.content.clone())); + params.push(SqlParam::Text(event.digest.clone())); + let sql: &str = match dialect { + Dialect::Sqlite => SQLITE_EVENT_INSERT, + Dialect::Postgres => POSTGRES_EVENT_INSERT, + }; + Query { + sql: sql.to_owned(), + params, + } +} diff --git a/crates/sqlbuild-state/src/execution_history/main/event_page_query.rs b/crates/sqlbuild-state/src/execution_history/main/event_page_query.rs new file mode 100644 index 0000000000..dea35f6dbd --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/event_page_query.rs @@ -0,0 +1,61 @@ +//! `event_page_query`: one ascending page of lifecycle events after a storage position. + +use crate::execution_history::_helpers::dialects::{event_table, placeholder, timestamp_param}; +use crate::execution_history::models::{Dialect, EventFilterSpec, Query, SqlParam}; + +/// `SELECT *` of events after `after_storage_id` matching the filter, one row past `limit`. +pub fn event_page_query( + dialect: Dialect, + after_storage_id: i64, + event_filter: &EventFilterSpec, + limit: i64, +) -> Query { + let mark: &str = placeholder(dialect); + let mut conditions: Vec = vec![format!("storage_id > {mark}")]; + let mut params: Vec = vec![SqlParam::Integer(after_storage_id)]; + for (column, value) in [ + ("invocation_id", &event_filter.invocation_id), + ("run_id", &event_filter.run_id), + ("producer", &event_filter.producer), + ] { + if let Some(value) = value { + conditions.push(format!("{column} = {mark}")); + params.push(SqlParam::Text(value.clone())); + } + } + if !event_filter.event_types.is_empty() { + match dialect { + Dialect::Sqlite => { + let marks: Vec<&str> = vec![mark; event_filter.event_types.len()]; + conditions.push(format!("event_type IN ({})", marks.join(","))); + params.extend(event_filter.event_types.iter().cloned().map(SqlParam::Text)); + } + Dialect::Postgres => { + conditions.push(format!("event_type = ANY({mark})")); + params.push(SqlParam::TextList(event_filter.event_types.clone())); + } + } + } + if let Some(family) = &event_filter.family { + conditions.push(format!("event_type LIKE {mark}")); + params.push(SqlParam::Text(format!("{family}_%"))); + } + for (comparison, bound) in [ + (">=", &event_filter.occurred_at_start), + ("<=", &event_filter.occurred_at_end), + ] { + if let Some(bound) = bound { + conditions.push(format!("occurred_at {comparison} {mark}")); + params.push(timestamp_param(dialect, &bound.instant)); + } + } + params.push(SqlParam::Integer(limit + 1)); + Query { + sql: format!( + "SELECT * FROM {} WHERE {} ORDER BY storage_id ASC LIMIT {mark}", + event_table(dialect), + conditions.join(" AND ") + ), + params, + } +} diff --git a/crates/sqlbuild-state/src/execution_history/main/mod.rs b/crates/sqlbuild-state/src/execution_history/main/mod.rs new file mode 100644 index 0000000000..91e45e99af --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/mod.rs @@ -0,0 +1,20 @@ +pub mod canonical_event_id; +pub mod decoded_run_cursor; +pub mod event_cursor; +pub mod event_cursor_position; +pub mod event_insert_query; +pub mod event_page_query; +pub mod prepared_events; +pub mod project_runs; +pub mod run_cursor; +pub mod run_insert_query; +pub mod run_page_query; +pub mod schema_migration_needed; +pub mod sqlite_database_path; +pub mod validate_filter; +pub mod validate_page; +pub mod validate_page_limit; +pub mod validate_postgres_settings; +pub mod validate_schema_target; +pub mod validate_stored_content; +pub mod validate_stored_event; diff --git a/crates/sqlbuild-state/src/execution_history/main/prepared_events.rs b/crates/sqlbuild-state/src/execution_history/main/prepared_events.rs new file mode 100644 index 0000000000..313239880e --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/prepared_events.rs @@ -0,0 +1,39 @@ +//! `prepared_events`: digest a batch and reject ids reused with different content inside it. + +use std::collections::HashMap; + +use sha2::{Digest, Sha256}; +use sqlbuild_core::text::main::python_str_repr::python_str_repr; + +use crate::execution_history::errors::{HistoryError, HistoryErrorKind}; +use crate::execution_history::models::{EventSubmission, PreparedEvent}; + +/// Each submission with the SHA-256 hex digest of its content, in order. +pub fn prepared_events( + submissions: Vec, +) -> Result, HistoryError> { + let mut digests: HashMap = HashMap::new(); + let mut prepared: Vec = Vec::with_capacity(submissions.len()); + for submission in submissions { + let digest: String = format!("{:x}", Sha256::digest(submission.content.as_bytes())); + if let Some(previous) = digests.get(&submission.event_id) + && *previous != digest + { + return Err(HistoryError::new( + HistoryErrorKind::IntegrityConflict, + format!( + "event_id {} has different canonical content within batch", + python_str_repr(&submission.event_id) + ), + )); + } + digests.insert(submission.event_id.clone(), digest.clone()); + prepared.push(PreparedEvent { + event_id: submission.event_id, + content: submission.content, + digest, + envelope: submission.envelope, + }); + } + Ok(prepared) +} diff --git a/crates/sqlbuild-state/src/execution_history/main/project_runs.rs b/crates/sqlbuild-state/src/execution_history/main/project_runs.rs new file mode 100644 index 0000000000..0d0978e3d0 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/project_runs.rs @@ -0,0 +1,111 @@ +//! `project_runs`: apply durable event facts in storage order to run projections. + +use std::collections::HashMap; + +use sqlbuild_core::text::main::python_str_repr::python_str_repr; + +use crate::execution_history::constants::{ + DUPLICATE_ORDER, RUN_COMPLETED_EVENT_TYPE, RUN_FAILED_EVENT_TYPE, RUN_STARTED_EVENT_TYPE, +}; +use crate::execution_history::errors::{HistoryError, HistoryErrorKind}; +use crate::execution_history::models::{ProjectionFact, RunFact, RunRow, RunStatus}; + +/// Facts sorted by storage order with identical duplicates dropped. +fn unique_ordered(facts: Vec) -> Result, HistoryError> { + let mut ordered: Vec = facts; + ordered.sort_by_key(|fact| fact.storage_order); + let mut unique: Vec = Vec::with_capacity(ordered.len()); + for fact in ordered { + if let Some(previous) = unique.last() + && previous.storage_order == fact.storage_order + { + if *previous == fact { + continue; + } + return Err(HistoryError::new( + HistoryErrorKind::ProjectionConsistency, + DUPLICATE_ORDER, + )); + } + unique.push(fact); + } + Ok(unique) +} + +/// The run after one of its events. +fn applied( + fact: &ProjectionFact, + run: &RunFact, + existing: Option, +) -> Result { + if let Some(existing) = &existing + && existing.invocation_id != run.invocation_id + { + return Err(HistoryError::new( + HistoryErrorKind::ProjectionConsistency, + format!( + "run_id {} has conflicting invocation correlations", + python_str_repr(&run.run_id) + ), + )); + } + let base: RunRow = existing.unwrap_or_else(|| RunRow { + run_id: run.run_id.clone(), + invocation_id: run.invocation_id.clone(), + created_at: fact.received_at, + status: RunStatus::Unknown, + is_complete: false, + last_event_cursor: fact.cursor.clone(), + last_storage_order: fact.storage_order, + command: None, + target: None, + environment: None, + started_at: None, + ended_at: None, + }); + let started_at = if run.event_type == RUN_STARTED_EVENT_TYPE { + Some(run.occurred_at) + } else { + base.started_at + }; + let (status, is_complete, ended_at) = match run.event_type.as_str() { + RUN_COMPLETED_EVENT_TYPE => (RunStatus::Completed, true, Some(run.occurred_at)), + RUN_FAILED_EVENT_TYPE => (RunStatus::Failed, true, Some(run.occurred_at)), + _ => (base.status, base.is_complete, base.ended_at), + }; + Ok(RunRow { + status, + is_complete, + started_at, + ended_at, + last_event_cursor: fact.cursor.clone(), + last_storage_order: fact.storage_order, + ..base + }) +} + +/// Current runs updated by every newer run-correlated fact, sorted by `(created_at, run_id)`. +pub fn project_runs( + facts: Vec, + current_runs: Vec, +) -> Result, HistoryError> { + let mut projected: HashMap = current_runs + .into_iter() + .map(|run| (run.run_id.clone(), run)) + .collect(); + for fact in unique_ordered(facts)? { + let Some(run) = &fact.run else { + continue; + }; + let next: RunRow = match projected.remove(&run.run_id) { + Some(existing) if fact.storage_order <= existing.last_storage_order => existing, + existing => applied(&fact, run, existing)?, + }; + projected.insert(run.run_id.clone(), next); + } + let mut runs: Vec = projected.into_values().collect(); + runs.sort_by(|left, right| { + (left.created_at, &left.run_id).cmp(&(right.created_at, &right.run_id)) + }); + Ok(runs) +} diff --git a/crates/sqlbuild-state/src/execution_history/main/run_cursor.rs b/crates/sqlbuild-state/src/execution_history/main/run_cursor.rs new file mode 100644 index 0000000000..c58423c4cf --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/run_cursor.rs @@ -0,0 +1,31 @@ +//! `run_cursor`: the opaque cursor of one projected run. + +use sqlbuild_core::json::main::dumps::dumps; +use sqlbuild_core::json::models::{JsonDialect, JsonValue, StdlibJsonOptions}; + +use crate::execution_history::_helpers::base64url::encode; +use crate::execution_history::_helpers::dialects::run_cursor_prefix; +use crate::execution_history::_helpers::timestamps::{default_text, stored_text}; +use crate::execution_history::errors::{HistoryError, HistoryErrorKind}; +use crate::execution_history::models::{Dialect, RunRow}; + +/// The prefix and URL-safe base64 of compact ASCII JSON `[namespace, created_at, run_id]`. +pub fn run_cursor(dialect: Dialect, namespace: &str, run: &RunRow) -> Result { + let created_at: String = match dialect { + Dialect::Sqlite => stored_text(&run.created_at), + Dialect::Postgres => default_text(&run.created_at), + }; + let payload = JsonValue::Array(vec![ + JsonValue::String(namespace.to_owned()), + JsonValue::String(created_at), + JsonValue::String(run.run_id.clone()), + ]); + let options = StdlibJsonOptions::new(None).with_separators(",", ":"); + let text: String = dumps(&payload, &JsonDialect::Stdlib(options)) + .map_err(|error| HistoryError::new(HistoryErrorKind::Storage, format!("{error:?}")))?; + Ok(format!( + "{}{}", + run_cursor_prefix(dialect), + encode(text.as_bytes()) + )) +} diff --git a/crates/sqlbuild-state/src/execution_history/main/run_insert_query.rs b/crates/sqlbuild-state/src/execution_history/main/run_insert_query.rs new file mode 100644 index 0000000000..d06d307a50 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/run_insert_query.rs @@ -0,0 +1,37 @@ +//! `run_insert_query`: the insert of one projected run. + +use crate::execution_history::_helpers::dialects::{ + optional_text_param, optional_timestamp_param, timestamp_param, +}; +use crate::execution_history::constants::{POSTGRES_RUN_INSERT, SCHEMA_VERSION, SQLITE_RUN_INSERT}; +use crate::execution_history::models::{Dialect, Query, RunRow, SqlParam}; + +/// SQLite stores `is_complete` as an integer, PostgreSQL as a boolean. +pub fn run_insert_query(dialect: Dialect, run: &RunRow) -> Query { + let is_complete: SqlParam = match dialect { + Dialect::Sqlite => SqlParam::Integer(i64::from(run.is_complete)), + Dialect::Postgres => SqlParam::Flag(run.is_complete), + }; + let sql: &str = match dialect { + Dialect::Sqlite => SQLITE_RUN_INSERT, + Dialect::Postgres => POSTGRES_RUN_INSERT, + }; + Query { + sql: sql.to_owned(), + params: vec![ + SqlParam::Text(run.run_id.clone()), + SqlParam::Text(run.invocation_id.clone()), + timestamp_param(dialect, &run.created_at), + SqlParam::Text(run.status.value().to_owned()), + is_complete, + SqlParam::Text(run.last_event_cursor.clone()), + SqlParam::Integer(run.last_storage_order), + optional_text_param(run.command.as_ref()), + optional_text_param(run.target.as_ref()), + optional_text_param(run.environment.as_ref()), + optional_timestamp_param(dialect, run.started_at.as_ref()), + optional_timestamp_param(dialect, run.ended_at.as_ref()), + SqlParam::Integer(SCHEMA_VERSION), + ], + } +} diff --git a/crates/sqlbuild-state/src/execution_history/main/run_page_query.rs b/crates/sqlbuild-state/src/execution_history/main/run_page_query.rs new file mode 100644 index 0000000000..09a6417515 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/run_page_query.rs @@ -0,0 +1,74 @@ +//! `run_page_query`: one page of projected runs after a `(created_at, run_id)` key. + +use crate::execution_history::_helpers::dialects::{placeholder, run_table, timestamp_param}; +use crate::execution_history::models::{Dialect, Query, RunCursorKey, RunFilterSpec, SqlParam}; + +/// `SELECT *` of runs after `after` matching the filter, one row past `limit`. +pub fn run_page_query( + dialect: Dialect, + after: Option<&RunCursorKey>, + run_filter: &RunFilterSpec, + limit: i64, +) -> Query { + let mark: &str = placeholder(dialect); + let mut conditions: Vec = Vec::new(); + let mut params: Vec = Vec::new(); + if let Some(after) = after { + match dialect { + Dialect::Sqlite => { + conditions.push("(created_at > ? OR (created_at = ? AND run_id > ?))".to_owned()); + params.push(SqlParam::Text(after.created_at_text.clone())); + params.push(SqlParam::Text(after.created_at_text.clone())); + } + Dialect::Postgres => { + conditions.push("(created_at, run_id) > (%s, %s)".to_owned()); + params.push(SqlParam::Timestamp(after.created_at)); + } + } + params.push(SqlParam::Text(after.run_id.clone())); + } + if let Some(invocation_id) = &run_filter.invocation_id { + conditions.push(format!("invocation_id = {mark}")); + params.push(SqlParam::Text(invocation_id.clone())); + } + if !run_filter.statuses.is_empty() { + let statuses: Vec = run_filter + .statuses + .iter() + .map(|status| status.value().to_owned()) + .collect(); + match dialect { + Dialect::Sqlite => { + let marks: Vec<&str> = vec![mark; statuses.len()]; + conditions.push(format!("status IN ({})", marks.join(","))); + params.extend(statuses.into_iter().map(SqlParam::Text)); + } + Dialect::Postgres => { + conditions.push(format!("status = ANY({mark})")); + params.push(SqlParam::TextList(statuses)); + } + } + } + for (comparison, bound) in [ + (">=", &run_filter.created_at_start), + ("<=", &run_filter.created_at_end), + ] { + if let Some(bound) = bound { + conditions.push(format!("created_at {comparison} {mark}")); + params.push(timestamp_param(dialect, &bound.instant)); + } + } + let filter: String = if conditions.is_empty() { + String::new() + } else { + format!(" WHERE {}", conditions.join(" AND ")) + }; + params.push(SqlParam::Integer(limit + 1)); + Query { + sql: format!( + "SELECT * FROM {}{filter} ORDER BY created_at, run_id LIMIT {mark}", + run_table(dialect) + ), + params, + } +} diff --git a/crates/sqlbuild-state/src/execution_history/main/schema_migration_needed.rs b/crates/sqlbuild-state/src/execution_history/main/schema_migration_needed.rs new file mode 100644 index 0000000000..be8f8e94aa --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/schema_migration_needed.rs @@ -0,0 +1,25 @@ +//! `schema_migration_needed`: whether a stored schema version still needs the current schema. + +use crate::execution_history::_helpers::schema::unsupported_version; +use crate::execution_history::constants::NO_SCHEMA; +use crate::execution_history::constants::SCHEMA_VERSION; +use crate::execution_history::errors::{HistoryError, HistoryErrorKind}; +use crate::execution_history::models::Dialect; + +/// `false` at the current version, `true` below it, an error when newer or negative. +pub fn schema_migration_needed(dialect: Dialect, current: i64) -> Result { + let backend: &str = match dialect { + Dialect::Sqlite => "SQLite history", + Dialect::Postgres => "PostgreSQL execution history", + }; + if current > SCHEMA_VERSION { + return Err(HistoryError::new( + HistoryErrorKind::UnsupportedSchemaVersion, + format!("{backend} schema {current} is newer than supported {SCHEMA_VERSION}"), + )); + } + if current < NO_SCHEMA { + return Err(unsupported_version(dialect, current)); + } + Ok(current != SCHEMA_VERSION) +} diff --git a/crates/sqlbuild-state/src/execution_history/main/sqlite_database_path.rs b/crates/sqlbuild-state/src/execution_history/main/sqlite_database_path.rs new file mode 100644 index 0000000000..29f794f727 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/sqlite_database_path.rs @@ -0,0 +1,42 @@ +//! `sqlite_database_path`: where a SQLite history lives, once its settings are valid. + +use sqlbuild_core::posix_path::main::posix_join::posix_join; +use sqlbuild_core::posix_path::main::posix_parts::posix_parts; + +use crate::execution_history::constants::{ + MEMORY_PATH, MIN_TIMEOUT, SQLITE_DIRECTORY_NAME, SQLITE_FILE_NAME, +}; +use crate::execution_history::errors::{HistoryError, HistoryErrorKind}; +use crate::execution_history::models::SqliteLocation; + +/// `:memory:`, the normalised explicit path, or `/.sqlbuild/history.sqlite3`; the busy timeout must be positive. +pub fn sqlite_database_path(location: &SqliteLocation) -> Result { + if location.project_dir.is_some() && location.path.is_some() { + return Err(HistoryError::new( + HistoryErrorKind::Storage, + "provide project_dir or path, not both", + )); + } + if location + .busy_timeout_ms + .is_none_or(|timeout| timeout < MIN_TIMEOUT) + { + return Err(HistoryError::new( + HistoryErrorKind::Storage, + "busy_timeout_ms must be a positive integer", + )); + } + let selected: &str = location + .path + .as_deref() + .or(location.project_dir.as_deref()) + .unwrap_or(&location.working_directory); + if selected == MEMORY_PATH { + return Ok(MEMORY_PATH.to_owned()); + } + let mut parts: Vec<&str> = posix_parts(selected); + if location.path.is_none() { + parts.extend([SQLITE_DIRECTORY_NAME, SQLITE_FILE_NAME]); + } + Ok(posix_join(&parts)) +} diff --git a/crates/sqlbuild-state/src/execution_history/main/validate_filter.rs b/crates/sqlbuild-state/src/execution_history/main/validate_filter.rs new file mode 100644 index 0000000000..2d662047af --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/validate_filter.rs @@ -0,0 +1,54 @@ +//! `validate_filter`: reject contradictory or malformed `EventFilter` and `RunFilter` fields. + +use crate::execution_history::_helpers::filters::{ + bounds_ordered, filled_text, filter_error, utc_bound, +}; +use crate::execution_history::errors::HistoryError; +use crate::execution_history::models::{EventFilterSpec, FilterSpec, RunFilterSpec}; + +/// Text fields non-blank, `event_types` exclusive with `family`, UTC ordered bounds. +fn validate_event_filter(event_filter: &EventFilterSpec) -> Result<(), HistoryError> { + filled_text(event_filter.invocation_id.as_deref(), "invocation_id")?; + filled_text(event_filter.run_id.as_deref(), "run_id")?; + filled_text(event_filter.producer.as_deref(), "producer")?; + if !event_filter.event_types.is_empty() && event_filter.family.is_some() { + return Err(filter_error( + "event_types and family are mutually exclusive", + )); + } + for event_type in &event_filter.event_types { + filled_text(Some(event_type), "event_types") + .map_err(|_| filter_error("event_types must contain only non-empty values"))?; + } + utc_bound(event_filter.occurred_at_start.as_ref(), "occurred_at_start")?; + utc_bound(event_filter.occurred_at_end.as_ref(), "occurred_at_end")?; + bounds_ordered( + ( + event_filter.occurred_at_start.as_ref(), + event_filter.occurred_at_end.as_ref(), + ), + "occurred_at_start must not be after occurred_at_end", + ) +} + +/// A non-blank `invocation_id` and UTC ordered creation bounds. +fn validate_run_filter(run_filter: &RunFilterSpec) -> Result<(), HistoryError> { + filled_text(run_filter.invocation_id.as_deref(), "invocation_id")?; + utc_bound(run_filter.created_at_start.as_ref(), "created_at_start")?; + utc_bound(run_filter.created_at_end.as_ref(), "created_at_end")?; + bounds_ordered( + ( + run_filter.created_at_start.as_ref(), + run_filter.created_at_end.as_ref(), + ), + "created_at_start must not be after created_at_end", + ) +} + +/// Reject the first invalid field of an event or run filter, in Python's check order. +pub fn validate_filter(filter: FilterSpec<'_>) -> Result<(), HistoryError> { + match filter { + FilterSpec::Event(event_filter) => validate_event_filter(event_filter), + FilterSpec::Run(run_filter) => validate_run_filter(run_filter), + } +} diff --git a/crates/sqlbuild-state/src/execution_history/main/validate_page.rs b/crates/sqlbuild-state/src/execution_history/main/validate_page.rs new file mode 100644 index 0000000000..d8c02771b2 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/validate_page.rs @@ -0,0 +1,25 @@ +//! `validate_page`: a page's next cursor must agree with its records. + +use crate::execution_history::errors::{HistoryError, HistoryErrorKind}; + +/// No cursor on an empty page; on an event page, the cursor of its last record. +pub fn validate_page( + is_event_page: bool, + has_records: bool, + (next_cursor, last_cursor): (Option<&str>, Option<&str>), +) -> Result<(), HistoryError> { + let kind: &str = if is_event_page { "event" } else { "run" }; + if !has_records && next_cursor.is_some() { + return Err(HistoryError::new( + HistoryErrorKind::ProjectionConsistency, + format!("an empty {kind} page cannot have a next cursor"), + )); + } + if is_event_page && has_records && next_cursor != last_cursor { + return Err(HistoryError::new( + HistoryErrorKind::ProjectionConsistency, + "event page cursor must identify its last record", + )); + } + Ok(()) +} diff --git a/crates/sqlbuild-state/src/execution_history/main/validate_page_limit.rs b/crates/sqlbuild-state/src/execution_history/main/validate_page_limit.rs new file mode 100644 index 0000000000..a030d8aad9 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/validate_page_limit.rs @@ -0,0 +1,16 @@ +//! `validate_page_limit`: a positive storage page limit no greater than the maximum. + +use crate::execution_history::constants::MAX_PAGE_LIMIT; +use crate::execution_history::constants::MIN_PAGE_LIMIT; +use crate::execution_history::errors::{HistoryError, HistoryErrorKind}; + +/// Reject a missing (non-integer) limit or one outside `1..=MAX_PAGE_LIMIT`. +pub fn validate_page_limit(limit: Option) -> Result<(), HistoryError> { + match limit { + Some(limit) if (MIN_PAGE_LIMIT..=MAX_PAGE_LIMIT).contains(&limit) => Ok(()), + _ => Err(HistoryError::new( + HistoryErrorKind::InvalidLimit, + format!("limit must be an integer from 1 through {MAX_PAGE_LIMIT}"), + )), + } +} diff --git a/crates/sqlbuild-state/src/execution_history/main/validate_postgres_settings.rs b/crates/sqlbuild-state/src/execution_history/main/validate_postgres_settings.rs new file mode 100644 index 0000000000..c102acbf86 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/validate_postgres_settings.rs @@ -0,0 +1,62 @@ +//! `validate_postgres_settings`: the PostgreSQL history DSN, timeouts and retry budget. + +use sqlbuild_core::text::main::python_strip::python_strip; + +use crate::execution_history::constants::{ + MIN_RETRIES, MIN_TIMEOUT, POSTGRES_MAX_CONNECT_TIMEOUT_SECONDS, + POSTGRES_MAX_STATEMENT_TIMEOUT_MS, POSTGRES_MAX_TRANSACTION_RETRIES, +}; +use crate::execution_history::errors::{HistoryError, HistoryErrorKind}; +use crate::execution_history::models::PostgresSettings; + +fn storage_error(message: String) -> HistoryError { + HistoryError::new(HistoryErrorKind::Storage, message) +} + +/// A non-blank DSN, positive bounded timeouts and `0..=10` transaction retries. +pub fn validate_postgres_settings(settings: &PostgresSettings) -> Result<(), HistoryError> { + if settings + .dsn + .as_deref() + .is_none_or(|dsn| python_strip(dsn).is_empty()) + { + return Err(storage_error( + "PostgreSQL execution history requires an explicit secret-resolved DSN".to_owned(), + )); + } + let timeouts = [ + ("connect_timeout_seconds", settings.connect_timeout_seconds), + ("statement_timeout_ms", settings.statement_timeout_ms), + ]; + for (name, value) in timeouts { + if value.is_none_or(|value| value < MIN_TIMEOUT) { + return Err(storage_error(format!("{name} must be a positive integer"))); + } + } + if settings + .connect_timeout_seconds + .is_some_and(|value| value > POSTGRES_MAX_CONNECT_TIMEOUT_SECONDS) + { + return Err(storage_error(format!( + "connect_timeout_seconds must not exceed {POSTGRES_MAX_CONNECT_TIMEOUT_SECONDS}" + ))); + } + if settings + .statement_timeout_ms + .is_some_and(|value| value > POSTGRES_MAX_STATEMENT_TIMEOUT_MS) + { + return Err(storage_error(format!( + "statement_timeout_ms must not exceed {POSTGRES_MAX_STATEMENT_TIMEOUT_MS}" + ))); + } + if settings + .max_transaction_retries + .is_none_or(|value| !(MIN_RETRIES..=POSTGRES_MAX_TRANSACTION_RETRIES).contains(&value)) + { + return Err(storage_error(format!( + "max_transaction_retries must be a non-negative integer no greater than \ + {POSTGRES_MAX_TRANSACTION_RETRIES}" + ))); + } + Ok(()) +} diff --git a/crates/sqlbuild-state/src/execution_history/main/validate_schema_target.rs b/crates/sqlbuild-state/src/execution_history/main/validate_schema_target.rs new file mode 100644 index 0000000000..1e61cf2baf --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/validate_schema_target.rs @@ -0,0 +1,14 @@ +//! `validate_schema_target`: only the current schema version can be upgraded to. + +use crate::execution_history::_helpers::schema::unsupported_version; +use crate::execution_history::constants::SCHEMA_VERSION; +use crate::execution_history::errors::HistoryError; +use crate::execution_history::models::Dialect; + +/// Reject any target other than the current schema version. +pub fn validate_schema_target(dialect: Dialect, target: i64) -> Result<(), HistoryError> { + if target == SCHEMA_VERSION { + return Ok(()); + } + Err(unsupported_version(dialect, target)) +} diff --git a/crates/sqlbuild-state/src/execution_history/main/validate_stored_content.rs b/crates/sqlbuild-state/src/execution_history/main/validate_stored_content.rs new file mode 100644 index 0000000000..fcb644c9ac --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/validate_stored_content.rs @@ -0,0 +1,25 @@ +//! `validate_stored_content`: an already stored event must hold the same canonical content. + +use sqlbuild_core::text::main::python_str_repr::python_str_repr; + +use crate::execution_history::errors::{HistoryError, HistoryErrorKind}; +use crate::execution_history::models::PreparedEvent; + +/// Reject a stored `(content_digest, payload_json)` that differs from the prepared event. +pub fn validate_stored_content( + prepared: &PreparedEvent, + (stored_digest, stored_payload): (Option<&str>, Option<&str>), +) -> Result<(), HistoryError> { + if stored_digest == Some(prepared.digest.as_str()) + && stored_payload == Some(prepared.content.as_str()) + { + return Ok(()); + } + Err(HistoryError::new( + HistoryErrorKind::IntegrityConflict, + format!( + "event_id {} has different canonical content", + python_str_repr(&prepared.event_id) + ), + )) +} diff --git a/crates/sqlbuild-state/src/execution_history/main/validate_stored_event.rs b/crates/sqlbuild-state/src/execution_history/main/validate_stored_event.rs new file mode 100644 index 0000000000..9f1e5d66fc --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/main/validate_stored_event.rs @@ -0,0 +1,28 @@ +//! `validate_stored_event`: a stored event's position, cursor and UTC receipt time. + +use crate::execution_history::constants::FIRST_POSITION; +use crate::execution_history::errors::{HistoryError, HistoryErrorKind}; + +/// A positive `storage_order`, a non-empty cursor and a `received_at` in `UTC`. +pub fn validate_stored_event( + storage_order: i64, + cursor: &str, + received_at_is_utc: bool, +) -> Result<(), HistoryError> { + let message: Option<&str> = if storage_order < FIRST_POSITION { + Some("storage_order must be a positive integer") + } else if cursor.is_empty() { + Some("event cursor must be non-empty") + } else if !received_at_is_utc { + Some("received_at must use UTC") + } else { + None + }; + match message { + Some(message) => Err(HistoryError::new( + HistoryErrorKind::ProjectionConsistency, + message, + )), + None => Ok(()), + } +} diff --git a/crates/sqlbuild-state/src/execution_history/mod.rs b/crates/sqlbuild-state/src/execution_history/mod.rs new file mode 100644 index 0000000000..2bce10c5ca --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/mod.rs @@ -0,0 +1,11 @@ +//! Execution history: append-only lifecycle event log and rebuildable run projection storage. + +pub(crate) mod _helpers; +pub mod constants; +pub mod errors; +pub mod main; +pub mod models; +pub mod types; + +#[cfg(test)] +mod tests; diff --git a/crates/sqlbuild-state/src/execution_history/models.rs b/crates/sqlbuild-state/src/execution_history/models.rs new file mode 100644 index 0000000000..7e2078d0c5 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/models.rs @@ -0,0 +1,178 @@ +//! Execution history facts, projected runs, filters and the parameterised SQL backends run. + +use crate::execution_history::types::UtcTime; + +/// Which backend's SQL placeholders, cursors and column types to use. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Dialect { + Sqlite, + Postgres, +} + +/// Latest terminal status derivable from durable run facts. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum RunStatus { + Unknown, + Completed, + Failed, +} + +impl RunStatus { + /// The `RunStatus` enum value text. + pub fn value(self) -> &'static str { + match self { + Self::Unknown => "unknown", + Self::Completed => "completed", + Self::Failed => "failed", + } + } +} + +/// Stable envelope fields stored next to one event's canonical content. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct EventEnvelope { + pub schema_version: i64, + /// Values of `ENVELOPE_TEXT_FIELDS`, in that order. + pub text_fields: [Option; 9], + pub occurred_at: Option, +} + +/// One event to append: its required id, canonical JSON content and envelope. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct EventSubmission { + pub event_id: String, + pub content: String, + pub envelope: EventEnvelope, +} + +/// An event ready to insert, with the SHA-256 digest of its content. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PreparedEvent { + pub event_id: String, + pub content: String, + pub digest: String, + pub envelope: EventEnvelope, +} + +/// The run-correlated part of a known lifecycle event. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RunFact { + pub run_id: String, + pub invocation_id: String, + pub event_type: String, + pub occurred_at: UtcTime, +} + +/// One durable event as projection reads it; opaque or run-less events carry no `run`. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ProjectionFact { + pub storage_order: i64, + pub cursor: String, + pub received_at: UtcTime, + pub run: Option, +} + +/// One projected run, as `RunRecord`. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RunRow { + pub run_id: String, + pub invocation_id: String, + pub created_at: UtcTime, + pub status: RunStatus, + pub is_complete: bool, + pub last_event_cursor: String, + pub last_storage_order: i64, + pub command: Option, + pub target: Option, + pub environment: Option, + pub started_at: Option, + pub ended_at: Option, +} + +/// A filter bound: its instant in UTC and whether its `tzinfo` is exactly `UTC`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct FilterTime { + pub instant: UtcTime, + pub is_utc: bool, +} + +/// `EventFilter` fields. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct EventFilterSpec { + pub invocation_id: Option, + pub run_id: Option, + pub producer: Option, + pub event_types: Vec, + /// The `EventFamily` value. + pub family: Option, + pub occurred_at_start: Option, + pub occurred_at_end: Option, +} + +/// `RunFilter` fields. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct RunFilterSpec { + pub invocation_id: Option, + pub statuses: Vec, + pub created_at_start: Option, + pub created_at_end: Option, +} + +/// An event or run filter to validate. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum FilterSpec<'a> { + Event(&'a EventFilterSpec), + Run(&'a RunFilterSpec), +} + +/// One SQL parameter; the backend glue renders timestamps for its driver. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum SqlParam { + Null, + Text(String), + Integer(i64), + Flag(bool), + Timestamp(UtcTime), + TextList(Vec), +} + +/// SQL text and its parameters in placeholder order. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Query { + pub sql: String, + pub params: Vec, +} + +/// A decoded run cursor's position: the stored `created_at` text and the run id. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct DecodedRunCursor { + pub created_at_text: String, + pub run_id: String, +} + +/// A validated run cursor's page position. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RunCursorKey { + pub created_at: UtcTime, + pub created_at_text: String, + pub run_id: String, +} + +/// `PostgresExecutionHistory` constructor settings; non-integer values are `None`. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PostgresSettings { + pub dsn: Option, + pub connect_timeout_seconds: Option, + pub statement_timeout_ms: Option, + pub max_transaction_retries: Option, +} + +/// Where a SQLite history lives: an explicit `path`, a `project_dir`, or the working directory. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SqliteLocation { + pub project_dir: Option, + pub path: Option, + pub working_directory: String, + /// `busy_timeout_ms`, or `None` when it is not an integer. + pub busy_timeout_ms: Option, +} diff --git a/crates/sqlbuild-state/src/execution_history/tests.rs b/crates/sqlbuild-state/src/execution_history/tests.rs new file mode 100644 index 0000000000..74f3f3f8e1 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/tests.rs @@ -0,0 +1,12 @@ +#[path = "tests/test_cursors.rs"] +mod cursors; +#[path = "tests/helpers.rs"] +mod helpers; +#[path = "tests/test_projection.rs"] +mod projection; +#[path = "tests/test_queries.rs"] +mod queries; +#[path = "tests/test_types.rs"] +mod test_types; +#[path = "tests/test_validation.rs"] +mod validation; diff --git a/crates/sqlbuild-state/src/execution_history/tests/helpers.rs b/crates/sqlbuild-state/src/execution_history/tests/helpers.rs new file mode 100644 index 0000000000..6a7784f493 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/tests/helpers.rs @@ -0,0 +1,65 @@ +use sqlbuild_cursor::algebra::models::CivilDate; + +use crate::execution_history::models::{FilterTime, ProjectionFact, RunFact, RunRow, RunStatus}; +use crate::execution_history::types::UtcTime; + +pub(crate) fn at(day: u32, microsecond: u32) -> UtcTime { + UtcTime { + date: CivilDate { + year: 2026, + month: 1, + day, + }, + hour: 0, + minute: 0, + second: 0, + microsecond, + } +} + +pub(crate) fn bound(day: u32, is_utc: bool) -> FilterTime { + FilterTime { + instant: at(day, 0), + is_utc, + } +} + +pub(crate) fn fact(storage_order: i64, run: Option<(&str, &str, &str)>) -> ProjectionFact { + ProjectionFact { + storage_order, + cursor: format!("c{storage_order}"), + received_at: at(u32::try_from(storage_order).unwrap_or(1), 0), + run: run.map(|(run_id, invocation_id, event_type)| RunFact { + run_id: run_id.to_owned(), + invocation_id: invocation_id.to_owned(), + event_type: event_type.to_owned(), + occurred_at: at(10, 5), + }), + } +} + +pub(crate) fn run_row(run_id: &str, created_day: u32, last: i64) -> RunRow { + RunRow { + run_id: run_id.to_owned(), + invocation_id: "i".to_owned(), + created_at: at(created_day, 0), + status: RunStatus::Unknown, + is_complete: false, + last_event_cursor: format!("c{last}"), + last_storage_order: last, + command: None, + target: None, + environment: None, + started_at: None, + ended_at: None, + } +} + +pub(crate) fn completed(run: RunRow) -> RunRow { + RunRow { + status: RunStatus::Completed, + is_complete: true, + ended_at: Some(at(10, 5)), + ..run + } +} diff --git a/crates/sqlbuild-state/src/execution_history/tests/test_cursors.rs b/crates/sqlbuild-state/src/execution_history/tests/test_cursors.rs new file mode 100644 index 0000000000..8c4894c757 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/tests/test_cursors.rs @@ -0,0 +1,128 @@ +use crate::execution_history::main::decoded_run_cursor::decoded_run_cursor; +use crate::execution_history::main::event_cursor_position::event_cursor_position; +use crate::execution_history::main::run_cursor::run_cursor; +use crate::execution_history::models::Dialect; +use crate::execution_history::tests::helpers::run_row; +use crate::execution_history::tests::test_types::{ + CursorRejectionTestCase, EventCursorTestCase, RunCursorTestCase, +}; + +#[test] +fn given_runs_when_encoding_cursors_then_python_base64_json_round_trips() { + let test_cases = [ + RunCursorTestCase { + description: "SQLite cursors carry microsecond text", + dialect: Dialect::Sqlite, + run: run_row("r1", 1, 1), + expected_cursor: "sqlite-run:WyJucyIsIjIwMjYtMDEtMDFUMDA6MDA6MDAuMDAwMDAwKzAwOjAwIiwicjEiXQ==", + expected_created_at: "2026-01-01T00:00:00.000000+00:00", + }, + RunCursorTestCase { + description: "PostgreSQL cursors carry isoformat text and escape non-ASCII", + dialect: Dialect::Postgres, + run: run_row("r\u{e9}", 1, 1), + expected_cursor: "postgres-run:WyJucyIsIjIwMjYtMDEtMDFUMDA6MDA6MDArMDA6MDAiLCJyXHUwMGU5Il0=", + expected_created_at: "2026-01-01T00:00:00+00:00", + }, + ]; + + for test_case in test_cases { + let cursor = run_cursor(test_case.dialect, "ns", &test_case.run); + let decoded = cursor + .clone() + .and_then(|cursor| decoded_run_cursor(test_case.dialect, &cursor, "ns")) + .map(|decoded| (decoded.created_at_text, decoded.run_id)); + assert_eq!( + (cursor, decoded), + ( + Ok(test_case.expected_cursor.to_owned()), + Ok(( + test_case.expected_created_at.to_owned(), + test_case.run.run_id.clone() + )) + ), + "{}", + test_case.description + ); + } +} + +#[test] +fn given_foreign_run_cursors_when_decoding_then_they_are_malformed() { + let test_cases = [ + CursorRejectionTestCase { + description: "other backend prefix", + dialect: Dialect::Postgres, + cursor: "sqlite-run:WyJucyIsIjIwMjYtMDEtMDFUMDA6MDA6MDAuMDAwMDAwKzAwOjAwIiwicjEiXQ==", + expected_message: "run cursor is malformed or belongs to another backend", + }, + CursorRejectionTestCase { + description: "bad base64", + dialect: Dialect::Sqlite, + cursor: "sqlite-run:@@@@", + expected_message: "run cursor is malformed or belongs to another backend", + }, + CursorRejectionTestCase { + description: "another namespace", + dialect: Dialect::Sqlite, + cursor: "sqlite-run:WyJvdGhlciIsIngiLCJyIl0=", + expected_message: "run cursor is malformed or belongs to another backend", + }, + ]; + + for test_case in test_cases { + assert_eq!( + decoded_run_cursor(test_case.dialect, test_case.cursor, "ns") + .map_err(|error| error.message), + Err(test_case.expected_message.to_owned()), + "{}", + test_case.description + ); + } +} + +#[test] +fn given_event_cursors_when_reading_positions_then_python_rules_apply() { + let test_cases = [ + EventCursorTestCase { + description: "no cursor starts at zero", + cursor: None, + expected_position: Ok(0), + }, + EventCursorTestCase { + description: "own cursor", + cursor: Some("sqlite-event:ns:7"), + expected_position: Ok(7), + }, + EventCursorTestCase { + description: "too few parts", + cursor: Some("sqlite-event"), + expected_position: Err("event cursor is malformed or belongs to another backend"), + }, + EventCursorTestCase { + description: "non-numeric position", + cursor: Some("sqlite-event:ns:x"), + expected_position: Err("event cursor is malformed or belongs to another backend"), + }, + EventCursorTestCase { + description: "another namespace", + cursor: Some("sqlite-event:other:7"), + expected_position: Err("event cursor is not valid for this storage"), + }, + EventCursorTestCase { + description: "position zero", + cursor: Some("sqlite-event:ns:0"), + expected_position: Err("event cursor is not valid for this storage"), + }, + ]; + + for test_case in test_cases { + assert_eq!( + event_cursor_position(Dialect::Sqlite, test_case.cursor, "ns") + .map_err(|error| error.message), + test_case.expected_position.map_err(str::to_owned), + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-state/src/execution_history/tests/test_projection.rs b/crates/sqlbuild-state/src/execution_history/tests/test_projection.rs new file mode 100644 index 0000000000..e4969326f1 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/tests/test_projection.rs @@ -0,0 +1,77 @@ +use crate::execution_history::errors::HistoryErrorKind; +use crate::execution_history::main::project_runs::project_runs; +use crate::execution_history::models::RunRow; +use crate::execution_history::tests::helpers::{at, completed, fact, run_row}; +use crate::execution_history::tests::test_types::ProjectionTestCase; + +#[test] +fn given_durable_facts_when_projecting_then_runs_follow_python_projection() { + let started = RunRow { + started_at: Some(at(10, 5)), + ..run_row("r1", 1, 1) + }; + let test_cases = [ + ProjectionTestCase { + description: "start then completion in storage order, regardless of input order", + facts: vec![ + fact(2, Some(("r1", "i", "run_completed"))), + fact(1, Some(("r1", "i", "run_started"))), + ], + current: vec![], + expected_runs: Ok(vec![completed(RunRow { + last_event_cursor: "c2".to_owned(), + last_storage_order: 2, + ..started.clone() + })]), + }, + ProjectionTestCase { + description: "stale and run-less facts are ignored", + facts: vec![fact(1, Some(("r1", "i", "run_failed"))), fact(3, None)], + current: vec![run_row("r1", 1, 2)], + expected_runs: Ok(vec![run_row("r1", 1, 2)]), + }, + ProjectionTestCase { + description: "identical duplicates collapse", + facts: vec![ + fact(1, Some(("r1", "i", "run_started"))), + fact(1, Some(("r1", "i", "run_started"))), + ], + current: vec![], + expected_runs: Ok(vec![started.clone()]), + }, + ProjectionTestCase { + description: "different facts sharing a storage order are inconsistent", + facts: vec![ + fact(1, Some(("r1", "i", "run_started"))), + fact(1, Some(("r2", "i", "run_started"))), + ], + current: vec![], + expected_runs: Err(HistoryErrorKind::ProjectionConsistency), + }, + ProjectionTestCase { + description: "a run cannot move to another invocation", + facts: vec![fact(3, Some(("r1", "other", "run_started")))], + current: vec![run_row("r1", 1, 2)], + expected_runs: Err(HistoryErrorKind::ProjectionConsistency), + }, + ProjectionTestCase { + description: "runs sort by creation time then run id", + facts: vec![], + current: vec![run_row("b", 2, 1), run_row("c", 1, 2), run_row("a", 2, 3)], + expected_runs: Ok(vec![ + run_row("c", 1, 2), + run_row("a", 2, 3), + run_row("b", 2, 1), + ]), + }, + ]; + + for test_case in test_cases { + assert_eq!( + project_runs(test_case.facts, test_case.current).map_err(|error| error.kind), + test_case.expected_runs, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-state/src/execution_history/tests/test_queries.rs b/crates/sqlbuild-state/src/execution_history/tests/test_queries.rs new file mode 100644 index 0000000000..a7593bb3c1 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/tests/test_queries.rs @@ -0,0 +1,68 @@ +use crate::execution_history::main::event_page_query::event_page_query; +use crate::execution_history::main::run_page_query::run_page_query; +use crate::execution_history::models::{ + Dialect, EventFilterSpec, RunCursorKey, RunFilterSpec, RunStatus, +}; +use crate::execution_history::tests::helpers::{at, bound}; +use crate::execution_history::tests::test_types::QueryTestCase; + +#[test] +fn given_filters_when_building_page_queries_then_sql_matches_each_backend() { + let event_filter = EventFilterSpec { + run_id: Some("r".to_owned()), + event_types: vec!["a".to_owned(), "b".to_owned()], + occurred_at_end: Some(bound(2, true)), + ..EventFilterSpec::default() + }; + let run_filter = RunFilterSpec { + statuses: vec![RunStatus::Failed, RunStatus::Completed], + created_at_start: Some(bound(1, true)), + ..RunFilterSpec::default() + }; + let key = RunCursorKey { + created_at: at(1, 0), + created_at_text: "2026-01-01T00:00:00.000000+00:00".to_owned(), + run_id: "r".to_owned(), + }; + let test_cases = [ + QueryTestCase { + description: "SQLite events expand IN placeholders", + query: event_page_query(Dialect::Sqlite, 4, &event_filter, 10), + expected_sql: "SELECT * FROM lifecycle_event_log WHERE storage_id > ? AND run_id = ? AND event_type IN (?,?) AND occurred_at <= ? ORDER BY storage_id ASC LIMIT ?", + expected_param_count: 6, + }, + QueryTestCase { + description: "PostgreSQL events bind one array", + query: event_page_query(Dialect::Postgres, 0, &event_filter, 10), + expected_sql: "SELECT * FROM sqlbuild_lifecycle_event_log WHERE storage_id > %s AND run_id = %s AND event_type = ANY(%s) AND occurred_at <= %s ORDER BY storage_id ASC LIMIT %s", + expected_param_count: 5, + }, + QueryTestCase { + description: "SQLite runs after a cursor compare stored text", + query: run_page_query(Dialect::Sqlite, Some(&key), &run_filter, 5), + expected_sql: "SELECT * FROM run_projection WHERE (created_at > ? OR (created_at = ? AND run_id > ?)) AND status IN (?,?) AND created_at >= ? ORDER BY created_at, run_id LIMIT ?", + expected_param_count: 7, + }, + QueryTestCase { + description: "PostgreSQL runs after a cursor compare a row value", + query: run_page_query(Dialect::Postgres, Some(&key), &run_filter, 5), + expected_sql: "SELECT * FROM sqlbuild_run_projection WHERE (created_at, run_id) > (%s, %s) AND status = ANY(%s) AND created_at >= %s ORDER BY created_at, run_id LIMIT %s", + expected_param_count: 5, + }, + QueryTestCase { + description: "unfiltered runs have no WHERE clause", + query: run_page_query(Dialect::Sqlite, None, &RunFilterSpec::default(), 5), + expected_sql: "SELECT * FROM run_projection ORDER BY created_at, run_id LIMIT ?", + expected_param_count: 1, + }, + ]; + + for test_case in test_cases { + assert_eq!( + (test_case.query.sql.as_str(), test_case.query.params.len()), + (test_case.expected_sql, test_case.expected_param_count), + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-state/src/execution_history/tests/test_types.rs b/crates/sqlbuild-state/src/execution_history/tests/test_types.rs new file mode 100644 index 0000000000..2e1e74c0ab --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/tests/test_types.rs @@ -0,0 +1,55 @@ +use crate::execution_history::errors::HistoryErrorKind; +use crate::execution_history::models::{ + Dialect, EventFilterSpec, ProjectionFact, Query, RunFilterSpec, RunRow, +}; + +pub(crate) struct ProjectionTestCase { + pub(crate) description: &'static str, + pub(crate) facts: Vec, + pub(crate) current: Vec, + pub(crate) expected_runs: Result, HistoryErrorKind>, +} + +pub(crate) struct QueryTestCase { + pub(crate) description: &'static str, + pub(crate) query: Query, + pub(crate) expected_sql: &'static str, + pub(crate) expected_param_count: usize, +} + +pub(crate) struct RunCursorTestCase { + pub(crate) description: &'static str, + pub(crate) dialect: Dialect, + pub(crate) run: RunRow, + pub(crate) expected_cursor: &'static str, + pub(crate) expected_created_at: &'static str, +} + +pub(crate) struct CursorRejectionTestCase { + pub(crate) description: &'static str, + pub(crate) dialect: Dialect, + pub(crate) cursor: &'static str, + pub(crate) expected_message: &'static str, +} + +pub(crate) struct EventCursorTestCase { + pub(crate) description: &'static str, + pub(crate) cursor: Option<&'static str>, + pub(crate) expected_position: Result, +} + +pub(crate) struct FilterTestCase { + pub(crate) description: &'static str, + pub(crate) event_filter: EventFilterSpec, + pub(crate) run_filter: RunFilterSpec, + pub(crate) expected_event_error: Option<&'static str>, + pub(crate) expected_run_error: Option<&'static str>, +} + +pub(crate) struct PathTestCase { + pub(crate) description: &'static str, + pub(crate) project_dir: Option<&'static str>, + pub(crate) path: Option<&'static str>, + pub(crate) busy_timeout_ms: Option, + pub(crate) expected_path: Result<&'static str, &'static str>, +} diff --git a/crates/sqlbuild-state/src/execution_history/tests/test_validation.rs b/crates/sqlbuild-state/src/execution_history/tests/test_validation.rs new file mode 100644 index 0000000000..549e688f9d --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/tests/test_validation.rs @@ -0,0 +1,143 @@ +use crate::execution_history::main::sqlite_database_path::sqlite_database_path; +use crate::execution_history::main::validate_filter::validate_filter; +use crate::execution_history::models::{ + EventFilterSpec, FilterSpec, RunFilterSpec, SqliteLocation, +}; +use crate::execution_history::tests::helpers::bound; +use crate::execution_history::tests::test_types::{FilterTestCase, PathTestCase}; + +#[test] +fn given_filters_when_validating_then_python_messages_are_returned() { + let test_cases = [ + FilterTestCase { + description: "empty filters are valid", + event_filter: EventFilterSpec::default(), + run_filter: RunFilterSpec::default(), + expected_event_error: None, + expected_run_error: None, + }, + FilterTestCase { + description: "Python whitespace counts as blank", + event_filter: EventFilterSpec { + producer: Some(" \u{1c}".to_owned()), + ..EventFilterSpec::default() + }, + run_filter: RunFilterSpec { + invocation_id: Some("\t".to_owned()), + ..RunFilterSpec::default() + }, + expected_event_error: Some("producer must be non-empty when provided"), + expected_run_error: Some("invocation_id must be non-empty when provided"), + }, + FilterTestCase { + description: "types and family conflict before blank types", + event_filter: EventFilterSpec { + event_types: vec![" ".to_owned()], + family: Some("run".to_owned()), + ..EventFilterSpec::default() + }, + run_filter: RunFilterSpec::default(), + expected_event_error: Some("event_types and family are mutually exclusive"), + expected_run_error: None, + }, + FilterTestCase { + description: "non-UTC and reversed bounds", + event_filter: EventFilterSpec { + occurred_at_start: Some(bound(1, false)), + ..EventFilterSpec::default() + }, + run_filter: RunFilterSpec { + created_at_start: Some(bound(2, true)), + created_at_end: Some(bound(1, true)), + ..RunFilterSpec::default() + }, + expected_event_error: Some("occurred_at_start must use UTC"), + expected_run_error: Some("created_at_start must not be after created_at_end"), + }, + ]; + + for test_case in test_cases { + assert_eq!( + ( + validate_filter(FilterSpec::Event(&test_case.event_filter)) + .err() + .map(|error| error.message), + validate_filter(FilterSpec::Run(&test_case.run_filter)) + .err() + .map(|error| error.message), + ), + ( + test_case.expected_event_error.map(str::to_owned), + test_case.expected_run_error.map(str::to_owned), + ), + "{}", + test_case.description + ); + } +} + +#[test] +fn given_locations_when_resolving_sqlite_path_then_pathlib_normalisation_applies() { + let test_cases = [ + PathTestCase { + description: "project directory gets the history file", + project_dir: Some("/work/project/"), + path: None, + busy_timeout_ms: Some(5000), + expected_path: Ok("/work/project/.sqlbuild/history.sqlite3"), + }, + PathTestCase { + description: "explicit path is normalised", + project_dir: None, + path: Some("./data//history.db"), + busy_timeout_ms: Some(5000), + expected_path: Ok("data/history.db"), + }, + PathTestCase { + description: "in-memory database", + project_dir: None, + path: Some(":memory:"), + busy_timeout_ms: Some(5000), + expected_path: Ok(":memory:"), + }, + PathTestCase { + description: "working directory by default", + project_dir: None, + path: None, + busy_timeout_ms: Some(5000), + expected_path: Ok("/cwd/.sqlbuild/history.sqlite3"), + }, + PathTestCase { + description: "both locations conflict", + project_dir: Some("a"), + path: Some("b"), + busy_timeout_ms: None, + expected_path: Err("provide project_dir or path, not both"), + }, + PathTestCase { + description: "non-positive busy timeout", + project_dir: None, + path: Some("b"), + busy_timeout_ms: Some(0), + expected_path: Err("busy_timeout_ms must be a positive integer"), + }, + ]; + + for test_case in test_cases { + let location = SqliteLocation { + project_dir: test_case.project_dir.map(str::to_owned), + path: test_case.path.map(str::to_owned), + working_directory: "/cwd".to_owned(), + busy_timeout_ms: test_case.busy_timeout_ms, + }; + assert_eq!( + sqlite_database_path(&location).map_err(|error| error.message), + test_case + .expected_path + .map(str::to_owned) + .map_err(str::to_owned), + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-state/src/execution_history/types.rs b/crates/sqlbuild-state/src/execution_history/types.rs new file mode 100644 index 0000000000..4c8fb40761 --- /dev/null +++ b/crates/sqlbuild-state/src/execution_history/types.rs @@ -0,0 +1,6 @@ +//! Execution history type aliases. + +use sqlbuild_cursor::algebra::models::LocalDateTime; + +/// A UTC instant as its wall time with microsecond precision. +pub type UtcTime = LocalDateTime; diff --git a/crates/sqlbuild-state/src/lib.rs b/crates/sqlbuild-state/src/lib.rs index d62c133a7e..3e6c8463ae 100644 --- a/crates/sqlbuild-state/src/lib.rs +++ b/crates/sqlbuild-state/src/lib.rs @@ -3,4 +3,5 @@ #![forbid(unsafe_code)] pub mod events; +pub mod execution_history; pub mod freshness; diff --git a/fensu.toml b/fensu.toml index 86eca090dd..bf472a55e7 100644 --- a/fensu.toml +++ b/fensu.toml @@ -78,6 +78,11 @@ rule = "FFR309" path = "src/sqlbuild/compiler/pipeline/__init__.py" reason = "Native bindings own the compile-and-plan pipeline; this package declares the pipeline models and keeps the custom materialization loader that imports project code." +[[targets.python.rule_exceptions]] +rule = "FFR309" +path = "src/sqlbuild/runtime/execution_history/__init__.py" +reason = "The native SQLite and PostgreSQL backends and sqlbuild-state own execution history storage, projection and validation; this package only declares the history models, types, constants and errors the public facades re-export." + [[targets.python.rule_exceptions]] rule = "FFR309" path = "src/sqlbuild/runtime/contracts/__init__.py" diff --git a/python-boundary.toml b/python-boundary.toml index cd96621e9b..ec3df3e61e 100644 --- a/python-boundary.toml +++ b/python-boundary.toml @@ -225,9 +225,29 @@ reason = "public-api" note = "Documented observability API used throughout the executor." [[stays]] -path = "runtime/execution_history" +path = "runtime/execution_history/__init__.py" reason = "public-api" -note = "Documented SQLite and Postgres execution-history backends." +note = "Package marker of the public execution history shapes below; storage, projection and validation are native (sqlbuild-state execution_history)." + +[[stays]] +path = "runtime/execution_history/models.py" +reason = "public-api" +note = "Rule 3 shapes: StoredEvent, filter, page and RunRecord dataclasses re-exported by sqlbuild.execution_history; __post_init__ delegates validation to native validate_history_* functions." + +[[stays]] +path = "runtime/execution_history/types.py" +reason = "public-api" +note = "Rule 3 shapes: EventFamily, RunStatus, cursor aliases and the storage protocols re-exported by sqlbuild.execution_history." + +[[stays]] +path = "runtime/execution_history/exceptions.py" +reason = "public-api" +note = "Rule 3 shapes: public storage errors re-exported by sqlbuild.execution_history and raised by the native backends." + +[[stays]] +path = "runtime/execution_history/constants.py" +reason = "public-api" +note = "Rule 3 shapes: page limit and schema version constants re-exported by sqlbuild.execution_history." [[stays]] path = "runtime/execution_limits" diff --git a/scripts/fensu_policy/constants.py b/scripts/fensu_policy/constants.py index 950db72806..0347c73e70 100644 --- a/scripts/fensu_policy/constants.py +++ b/scripts/fensu_policy/constants.py @@ -194,9 +194,6 @@ STRICT_ADAPTER_CLASS_NAME: str = "StrictAdapter" ATTRIBUTE_REFERENCE_KIND: str = "attribute" -NON_RECURSIVE_EXECUTION_HISTORY_PATH: str = ( - "src/sqlbuild/runtime/execution_history/classes/postgres_execution_history.py" -) APPROVED_RAW_STATEMENT_MODULES: frozenset[str] = frozenset( { "src/sqlbuild/adapter/contract/classes/observed_connection.py", diff --git a/scripts/fensu_policy/rules/execution_observability.py b/scripts/fensu_policy/rules/execution_observability.py index 84ee4fd737..ccb892246d 100644 --- a/scripts/fensu_policy/rules/execution_observability.py +++ b/scripts/fensu_policy/rules/execution_observability.py @@ -42,7 +42,6 @@ LIFECYCLE_EVENT_CLASS_NAME, LIFECYCLE_EVENT_FACTORY_NAME, LIFECYCLE_FACTORY_OWNER_PATHS, - NON_RECURSIVE_EXECUTION_HISTORY_PATH, OBSERVABILITY_DEFINITION_MODULES, OBSERVABILITY_FACTORY_MODULES, OBSERVABILITY_LIFECYCLE_EVENT_MODULES, @@ -97,7 +96,7 @@ def raw_driver_execution_boundary(*, module: object, ctx: RuleContext) -> list[F if ctx.scope() != ROOT_SCOPE_NAME or ctx.repo_relative_parts()[:2] != RUNTIME_ROOT_PARTS: return [] path: str = "/".join(ctx.repo_relative_parts()) - if path == NON_RECURSIVE_EXECUTION_HISTORY_PATH or path in APPROVED_RAW_STATEMENT_MODULES: + if path in APPROVED_RAW_STATEMENT_MODULES: return [] tree: ast.Module = module if isinstance(module, ast.Module) else ast.parse(ctx.text.source) return [ diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index 362c97f69e..6afbf93f48 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -353,6 +353,19 @@ from sqlbuild.runtime.compute_logs.models import ( from sqlbuild.runtime.compute_logs.types import ComputeLogStream from sqlbuild.runtime.contracts.models import ConnectionHooks from sqlbuild.runtime.contracts.types import ExecutionResourceKind +from sqlbuild.runtime.execution_history.models import ( + EventFilter, + EventPage, + RunFilter, + RunPage, + RunRecord, + StoredEvent, +) +from sqlbuild.runtime.execution_history.types import ( + CanonicalLifecycleEvent, + LifecycleEventLogStorage, + RunStorage, +) from sqlbuild.runtime.observability.classes.event_dispatcher import EventDispatcher from sqlbuild.runtime.observability.models import DiagnosticLog, LifecycleEvent from sqlbuild.runtime.output_capture.models import ( @@ -3536,6 +3549,100 @@ class LocalFilesystemComputeLogStorage: traceback: TracebackType | None, ) -> bool | None: ... +class _ExecutionHistoryStorage: + @property + def project_calls(self) -> int: ... + def append_event(self, event: CanonicalLifecycleEvent) -> StoredEvent: ... + def append_events( + self, events: Iterable[CanonicalLifecycleEvent] + ) -> tuple[StoredEvent, ...]: ... + def get_events( + self, + *, + event_filter: EventFilter, + after_cursor: str | None = None, + limit: int = 100, + ) -> EventPage: ... + def get_run(self, run_id: str) -> RunRecord | None: ... + def get_runs( + self, + *, + run_filter: RunFilter, + after_cursor: str | None = None, + limit: int = 100, + ) -> RunPage: ... + def project(self, stored_events: Iterable[StoredEvent]) -> tuple[RunRecord, ...]: ... + def rebuild_from_events( + self, stored_events: Iterable[StoredEvent] + ) -> tuple[RunRecord, ...]: ... + def append_and_project( + self, events: Iterable[CanonicalLifecycleEvent] + ) -> tuple[StoredEvent, ...]: ... + def reconcile(self) -> tuple[RunRecord, ...]: ... + def check_health(self) -> bool: ... + def get_schema_version(self) -> int: ... + def upgrade_schema(self, *, target_version: int | None = None) -> int: ... + def close(self) -> None: ... + def dispose(self) -> None: ... + def __enter__(self) -> Self: ... + def __exit__( + self, + exception_type: type[BaseException] | None, + exception: BaseException | None, + traceback: TracebackType | None, + ) -> bool | None: ... + +class SQLiteExecutionHistory(_ExecutionHistoryStorage): + """Project-local SQLite implementation of event-log and run-projection storage contracts.""" + + _connection: Any + def __init__( + self, + project_dir: Path | str | None = None, + *, + path: Path | str | None = None, + busy_timeout_ms: int = 5_000, + ) -> None: ... + @property + def path(self) -> Path | None: ... + def _migrate(self) -> None: ... + def _read_all_events(self) -> tuple[StoredEvent, ...]: ... + def _publish_projection(self, projected: tuple[RunRecord, ...]) -> None: ... + +class PostgresExecutionHistory(_ExecutionHistoryStorage): + """Deployed PostgreSQL implementation of event-log and run-projection contracts.""" + + _connection: Any + def __init__( + self, + dsn: str, + *, + connect_timeout_seconds: int = 10, + statement_timeout_ms: int = 30_000, + max_transaction_retries: int = 3, + ) -> None: ... + def _append_events( + self, *, cursor: Any, events: tuple[CanonicalLifecycleEvent, ...] + ) -> tuple[StoredEvent, ...]: ... + def _replace_projection(self, *, cursor: Any, projected: tuple[RunRecord, ...]) -> None: ... + def _read_all_events(self, *, cursor: Any) -> tuple[StoredEvent, ...]: ... + +def canonical_event_id(event: CanonicalLifecycleEvent) -> str: ... +def canonical_event_content(event: CanonicalLifecycleEvent) -> str: ... +def project_execution_runs( + *, stored_events: Iterable[StoredEvent], current_runs: Iterable[RunRecord] | None = None +) -> tuple[RunRecord, ...]: ... +def validate_history_page_limit(limit: int) -> None: ... +def append_and_project( + *, + event_log: LifecycleEventLogStorage, + run_storage: RunStorage, + events: Iterable[CanonicalLifecycleEvent], +) -> tuple[StoredEvent, ...]: ... +def validate_stored_history_event(record: StoredEvent) -> None: ... +def validate_history_event_filter(event_filter: EventFilter) -> None: ... +def validate_history_run_filter(run_filter: RunFilter) -> None: ... +def validate_history_page(*, page: EventPage | RunPage, is_event_page: bool) -> None: ... def current_native_progress_projector() -> NativeProgressProjector | None: ... class ConnectionProgressReporter: diff --git a/src/sqlbuild/execution_history.py b/src/sqlbuild/execution_history.py index 39f5217e07..a3674fdda6 100644 --- a/src/sqlbuild/execution_history.py +++ b/src/sqlbuild/execution_history.py @@ -2,6 +2,11 @@ from collections.abc import Iterable +from sqlbuild._native import append_and_project as _append_and_project +from sqlbuild._native import canonical_event_content as _canonical_event_content +from sqlbuild._native import canonical_event_id as _canonical_event_id +from sqlbuild._native import project_execution_runs as _project_execution_runs +from sqlbuild._native import validate_history_page_limit as _validate_history_page_limit from sqlbuild.runtime.execution_history.constants import DEFAULT_PAGE_LIMIT, MAX_PAGE_LIMIT from sqlbuild.runtime.execution_history.exceptions import ( ExecutionHistoryStorageError, @@ -13,19 +18,6 @@ ProjectionConsistencyError, UnsupportedSchemaVersionError, ) -from sqlbuild.runtime.execution_history.main.append_and_project import ( - append_and_project as _append_and_project, -) -from sqlbuild.runtime.execution_history.main.canonical_event_content import ( - canonical_event_content as _canonical_event_content, -) -from sqlbuild.runtime.execution_history.main.canonical_event_id import ( - canonical_event_id as _canonical_event_id, -) -from sqlbuild.runtime.execution_history.main.project_runs import project_runs as _project_runs -from sqlbuild.runtime.execution_history.main.validate_page_limit import ( - validate_page_limit as _validate_page_limit, -) from sqlbuild.runtime.execution_history.models import ( EventFilter, EventPage, @@ -104,10 +96,10 @@ def project_runs( ) -> tuple[RunRecord, ...]: """Apply durable event facts in storage order to immutable run projections.""" - return _project_runs(stored_events=stored_events, current_runs=current_runs) + return _project_execution_runs(stored_events=stored_events, current_runs=current_runs) def validate_page_limit(limit: int) -> None: """Validate a positive page limit no greater than the public maximum.""" - _validate_page_limit(limit) + _validate_history_page_limit(limit) diff --git a/src/sqlbuild/postgres_history.py b/src/sqlbuild/postgres_history.py index 7c7d02d2a5..a26744c608 100644 --- a/src/sqlbuild/postgres_history.py +++ b/src/sqlbuild/postgres_history.py @@ -1,11 +1,9 @@ """Typed public PostgreSQL execution history API.""" -from sqlbuild.runtime.execution_history.classes.postgres_execution_history import ( - PostgresExecutionHistory as _PostgresExecutionHistory, -) +from sqlbuild._native import PostgresExecutionHistory as _NativePostgresExecutionHistory __all__ = ("PostgresExecutionHistory",) -class PostgresExecutionHistory(_PostgresExecutionHistory): +class PostgresExecutionHistory(_NativePostgresExecutionHistory): """Supported deployed PostgreSQL execution history backend.""" diff --git a/src/sqlbuild/runtime/execution_history/_helpers/__init__.py b/src/sqlbuild/runtime/execution_history/_helpers/__init__.py deleted file mode 100644 index 4ec9e74a57..0000000000 --- a/src/sqlbuild/runtime/execution_history/_helpers/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Private execution history helpers.""" diff --git a/src/sqlbuild/runtime/execution_history/_helpers/canonical_events.py b/src/sqlbuild/runtime/execution_history/_helpers/canonical_events.py deleted file mode 100644 index 7d636e5c1b..0000000000 --- a/src/sqlbuild/runtime/execution_history/_helpers/canonical_events.py +++ /dev/null @@ -1,24 +0,0 @@ -"""Canonical lifecycle event identity and content helpers.""" - -from sqlbuild.runtime.execution_history.exceptions import InvalidEventError -from sqlbuild.runtime.execution_history.types import CanonicalLifecycleEvent -from sqlbuild.runtime.observability.main.lifecycle_event_to_json import lifecycle_event_to_json -from sqlbuild.runtime.observability.models import LifecycleEvent - - -def canonical_event_id(*, event: CanonicalLifecycleEvent) -> str: - """Return the required stable event ID from a known or opaque envelope.""" - - event_id: object = ( - event.event_id if isinstance(event, LifecycleEvent) else event.raw.get("event_id") - ) - if not isinstance(event_id, str) or not event_id: - raise InvalidEventError("canonical lifecycle event requires a non-empty string event_id") - return event_id - - -def canonical_event_content(*, event: CanonicalLifecycleEvent) -> str: - """Return deterministic canonical content for idempotency comparisons.""" - - _ = canonical_event_id(event=event) - return lifecycle_event_to_json(event) diff --git a/src/sqlbuild/runtime/execution_history/_helpers/projection.py b/src/sqlbuild/runtime/execution_history/_helpers/projection.py deleted file mode 100644 index 3f9c076892..0000000000 --- a/src/sqlbuild/runtime/execution_history/_helpers/projection.py +++ /dev/null @@ -1,94 +0,0 @@ -"""Deterministic run projection from durable event facts.""" - -from collections.abc import Iterable -from dataclasses import replace -from datetime import datetime - -from sqlbuild.runtime.execution_history.constants import ( - RUN_COMPLETED_EVENT_TYPE, - RUN_FAILED_EVENT_TYPE, - RUN_STARTED_EVENT_TYPE, -) -from sqlbuild.runtime.execution_history.exceptions import ProjectionConsistencyError -from sqlbuild.runtime.execution_history.models import RunRecord, StoredEvent -from sqlbuild.runtime.execution_history.types import RunStatus -from sqlbuild.runtime.observability.models import LifecycleEvent - - -def project_runs( - *, stored_events: Iterable[StoredEvent], current_runs: Iterable[RunRecord] = () -) -> tuple[RunRecord, ...]: - """Apply durable events in storage order and return deterministic run projections.""" - - projected: dict[str, RunRecord] = {run.run_id: run for run in current_runs} - ordered_events: tuple[StoredEvent, ...] = _get_unique_ordered_events( - stored_events=stored_events - ) - for stored_event in ordered_events: - event: object = stored_event.event - if not isinstance(event, LifecycleEvent) or event.run_id is None: - continue - existing: RunRecord | None = projected.get(event.run_id) - if existing is not None and stored_event.storage_order <= existing.last_storage_order: - continue - projected[event.run_id] = _apply_run_event( - stored_event=stored_event, event=event, existing=existing - ) - return tuple(sorted(projected.values(), key=lambda run: (run.created_at, run.run_id))) - - -def _get_unique_ordered_events(*, stored_events: Iterable[StoredEvent]) -> tuple[StoredEvent, ...]: - ordered_events: tuple[StoredEvent, ...] = tuple( - sorted(stored_events, key=lambda stored_event: stored_event.storage_order) - ) - unique_events: list[StoredEvent] = [] - for stored_event in ordered_events: - previous: StoredEvent | None = unique_events[-1] if unique_events else None - if previous is not None and stored_event.storage_order == previous.storage_order: - if stored_event == previous: - continue - raise ProjectionConsistencyError("storage_order must uniquely identify a durable event") - unique_events.append(stored_event) - return tuple(unique_events) - - -def _apply_run_event( - *, stored_event: StoredEvent, event: LifecycleEvent, existing: RunRecord | None -) -> RunRecord: - if existing is not None and existing.invocation_id != event.invocation_id: - raise ProjectionConsistencyError( - f"run_id {event.run_id!r} has conflicting invocation correlations" - ) - if existing is None: - existing = RunRecord( - run_id=event.run_id or "", - invocation_id=event.invocation_id, - created_at=stored_event.received_at, - status=RunStatus.UNKNOWN, - is_complete=False, - last_event_cursor=stored_event.cursor, - last_storage_order=stored_event.storage_order, - ) - started_at: datetime | None = ( - event.occurred_at if event.event_type == RUN_STARTED_EVENT_TYPE else existing.started_at - ) - status, is_complete, ended_at = _terminal_state(event=event, existing=existing) - return replace( - existing, - status=status, - is_complete=is_complete, - started_at=started_at, - ended_at=ended_at, - last_event_cursor=stored_event.cursor, - last_storage_order=stored_event.storage_order, - ) - - -def _terminal_state( - *, event: LifecycleEvent, existing: RunRecord -) -> tuple[RunStatus, bool, datetime | None]: - if event.event_type == RUN_COMPLETED_EVENT_TYPE: - return RunStatus.COMPLETED, True, event.occurred_at - if event.event_type == RUN_FAILED_EVENT_TYPE: - return RunStatus.FAILED, True, event.occurred_at - return existing.status, existing.is_complete, existing.ended_at diff --git a/src/sqlbuild/runtime/execution_history/_helpers/validation.py b/src/sqlbuild/runtime/execution_history/_helpers/validation.py deleted file mode 100644 index ff0bfdcbf6..0000000000 --- a/src/sqlbuild/runtime/execution_history/_helpers/validation.py +++ /dev/null @@ -1,38 +0,0 @@ -"""Execution history contract validation helpers.""" - -from datetime import UTC, datetime - -from sqlbuild.runtime.execution_history.constants import MAX_PAGE_LIMIT -from sqlbuild.runtime.execution_history.exceptions import ( - InvalidFilterError, - InvalidLimitError, - ProjectionConsistencyError, -) - - -def validate_page_limit(limit: int) -> None: - """Validate a positive, bounded storage page limit.""" - - if isinstance(limit, bool) or not isinstance(limit, int) or not 1 <= limit <= MAX_PAGE_LIMIT: - raise InvalidLimitError(f"limit must be an integer from 1 through {MAX_PAGE_LIMIT}") - - -def validate_filter_text(*, value: str | None, field_name: str) -> None: - """Validate an optional non-empty filter correlation value.""" - - if value is not None and not value.strip(): - raise InvalidFilterError(f"{field_name} must be non-empty when provided") - - -def validate_filter_timestamp(*, value: datetime | None, field_name: str) -> None: - """Validate an optional UTC filter boundary.""" - - if value is not None and value.tzinfo is not UTC: - raise InvalidFilterError(f"{field_name} must use UTC") - - -def validate_storage_timestamp(*, value: datetime, field_name: str) -> None: - """Validate a UTC timestamp carried by a durable storage fact.""" - - if value.tzinfo is not UTC: - raise ProjectionConsistencyError(f"{field_name} must use UTC") diff --git a/src/sqlbuild/runtime/execution_history/classes/__init__.py b/src/sqlbuild/runtime/execution_history/classes/__init__.py deleted file mode 100644 index 9cefe90b49..0000000000 --- a/src/sqlbuild/runtime/execution_history/classes/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Execution history storage implementations.""" diff --git a/src/sqlbuild/runtime/execution_history/classes/postgres_execution_history.py b/src/sqlbuild/runtime/execution_history/classes/postgres_execution_history.py deleted file mode 100644 index 10dc04925b..0000000000 --- a/src/sqlbuild/runtime/execution_history/classes/postgres_execution_history.py +++ /dev/null @@ -1,868 +0,0 @@ -"""Deployed PostgreSQL execution history storage.""" - -from __future__ import annotations - -import base64 -import binascii -import hashlib -import json -from collections.abc import Callable, Iterable, Mapping -from datetime import UTC, datetime, timedelta -from types import TracebackType -from typing import Any, Self -from uuid import uuid4 - -from sqlbuild.runtime.execution_history.constants import ( - CURRENT_RUN_STORAGE_SCHEMA_VERSION, - DEFAULT_PAGE_LIMIT, -) -from sqlbuild.runtime.execution_history.exceptions import ( - ExecutionHistoryStorageError, - IntegrityConflictError, - InvalidCursorError, - UnsupportedSchemaVersionError, -) -from sqlbuild.runtime.execution_history.main.canonical_event_content import canonical_event_content -from sqlbuild.runtime.execution_history.main.canonical_event_id import canonical_event_id -from sqlbuild.runtime.execution_history.main.project_runs import project_runs -from sqlbuild.runtime.execution_history.main.validate_page_limit import validate_page_limit -from sqlbuild.runtime.execution_history.models import ( - EventFilter, - EventPage, - RunFilter, - RunPage, - RunRecord, - StoredEvent, -) -from sqlbuild.runtime.execution_history.types import CanonicalLifecycleEvent, RunStatus -from sqlbuild.runtime.observability.main.lifecycle_event_from_json import lifecycle_event_from_json -from sqlbuild.runtime.observability.models import LifecycleEvent, OpaqueLifecycleEvent - -_CONNECT_TIMEOUT_SECONDS: int = 10 -_STATEMENT_TIMEOUT_MS: int = 30_000 -_MAX_TRANSACTION_RETRIES: int = 3 -_MAX_CONNECT_TIMEOUT_SECONDS: int = 300 -_MAX_STATEMENT_TIMEOUT_MS: int = 3_600_000 -_MAX_TRANSACTION_RETRY_LIMIT: int = 10 -_SCHEMA_VERSION: int = 1 -_MIGRATION_LOCK_ID: int = 7_218_786_101 -_HISTORY_LOCK_ID: int = 7_218_786_102 -_EVENT_CURSOR_PREFIX: str = "postgres-event:" -_RUN_CURSOR_PREFIX: str = "postgres-run:" -_RUN_CURSOR_PART_COUNT: int = 3 -_RETRYABLE_SQLSTATES: frozenset[str] = frozenset(("40001", "40P01")) - -_CREATE_SCHEMA: tuple[str, ...] = ( - """CREATE TABLE IF NOT EXISTS sqlbuild_storage_migrations ( - singleton BOOLEAN PRIMARY KEY DEFAULT TRUE CHECK (singleton), - schema_version INTEGER NOT NULL, - storage_namespace TEXT NOT NULL, - migrated_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP - )""", - """CREATE TABLE IF NOT EXISTS sqlbuild_lifecycle_event_log ( - storage_id BIGINT GENERATED ALWAYS AS IDENTITY PRIMARY KEY, - event_id TEXT NOT NULL UNIQUE, - schema_version INTEGER NOT NULL, - producer TEXT, - producer_version TEXT, - event_type TEXT, - occurred_at TIMESTAMPTZ, - received_at TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP, - invocation_id TEXT, - run_id TEXT, - resource_id TEXT, - resource_attempt_id TEXT, - operation_id TEXT, - statement_id TEXT, - payload_json TEXT NOT NULL, - content_digest TEXT NOT NULL - )""", - """CREATE INDEX IF NOT EXISTS sqlbuild_lifecycle_event_log_run_storage - ON sqlbuild_lifecycle_event_log (run_id, storage_id)""", - """CREATE INDEX IF NOT EXISTS sqlbuild_lifecycle_event_log_invocation_storage - ON sqlbuild_lifecycle_event_log (invocation_id, storage_id)""", - """CREATE INDEX IF NOT EXISTS sqlbuild_lifecycle_event_log_type_storage - ON sqlbuild_lifecycle_event_log (event_type, storage_id)""", - """CREATE TABLE IF NOT EXISTS sqlbuild_run_projection ( - run_id TEXT PRIMARY KEY, - invocation_id TEXT NOT NULL, - created_at TIMESTAMPTZ NOT NULL, - status TEXT NOT NULL, - is_complete BOOLEAN NOT NULL, - last_event_cursor TEXT NOT NULL, - last_event_storage_id BIGINT NOT NULL, - command TEXT, - target TEXT, - environment TEXT, - started_at TIMESTAMPTZ, - ended_at TIMESTAMPTZ, - projection_schema_version INTEGER NOT NULL - )""", - """CREATE INDEX IF NOT EXISTS sqlbuild_run_projection_created - ON sqlbuild_run_projection (created_at, run_id)""", - """CREATE INDEX IF NOT EXISTS sqlbuild_run_projection_invocation_created - ON sqlbuild_run_projection (invocation_id, created_at, run_id)""", -) - - -class PostgresExecutionHistory: - """PostgreSQL implementation of event-log and run-projection contracts.""" - - def __init__( - self, - dsn: str, - *, - connect_timeout_seconds: int = _CONNECT_TIMEOUT_SECONDS, - statement_timeout_ms: int = _STATEMENT_TIMEOUT_MS, - max_transaction_retries: int = _MAX_TRANSACTION_RETRIES, - ) -> None: - self._validate_configuration( - dsn=dsn, - connect_timeout_seconds=connect_timeout_seconds, - statement_timeout_ms=statement_timeout_ms, - max_transaction_retries=max_transaction_retries, - ) - psycopg, dict_row = self._load_driver() - self._psycopg: Any = psycopg - self._max_transaction_retries = max_transaction_retries - self._closed = False - self._project_calls = 0 - connection: Any | None = None - try: - connection = psycopg.connect( - dsn, - connect_timeout=connect_timeout_seconds, - autocommit=True, - row_factory=dict_row, - ) - self._connection: Any = connection - with connection.cursor() as cursor: - cursor.execute( - "SELECT set_config('statement_timeout', %s, false)", - (str(statement_timeout_ms),), - ) - self._statement_timeout_ms = statement_timeout_ms - self._migrate() - _ = self.reconcile() - except ExecutionHistoryStorageError: - self._close_failed_connection(connection=connection) - self._closed = True - raise - except Exception: - self._close_failed_connection(connection=connection) - self._closed = True - raise ExecutionHistoryStorageError( - "cannot connect to PostgreSQL execution history; verify the secret-resolved DSN, " - "network access, and sqlbuild[postgres] installation" - ) from None - - def __repr__(self) -> str: - state: str = "closed" if self._closed else "open" - return f"{type(self).__name__}(state={state!r})" - - @property - def project_calls(self) -> int: - """Return the number of incremental projection calls.""" - - return self._project_calls - - def append_event(self, event: CanonicalLifecycleEvent) -> StoredEvent: - self._ensure_open() - return self.append_events((event,))[0] - - def append_events(self, events: Iterable[CanonicalLifecycleEvent]) -> tuple[StoredEvent, ...]: - self._ensure_open() - pending: tuple[CanonicalLifecycleEvent, ...] = tuple(events) - - def append(cursor: Any) -> tuple[StoredEvent, ...]: - self._lock_history(cursor=cursor) - return self._append_events(cursor=cursor, events=pending) - - return self._transactional(action=append, failure_message="PostgreSQL event append failed") - - def get_events( - self, - *, - event_filter: EventFilter, - after_cursor: str | None = None, - limit: int = DEFAULT_PAGE_LIMIT, - ) -> EventPage: - self._ensure_open() - validate_page_limit(limit) - after_storage_id: int = self._event_cursor_position(after_cursor) - conditions: list[str] = ["storage_id > %s"] - values: list[object] = [after_storage_id] - conditions, values = self._add_event_filters( - conditions=conditions, values=values, event_filter=event_filter - ) - query: str = ( - "SELECT * FROM sqlbuild_lifecycle_event_log WHERE " - + " AND ".join(conditions) - + " ORDER BY storage_id ASC LIMIT %s" - ) - values.append(limit + 1) - rows: list[Mapping[str, object]] = self._read_rows(query=query, values=values) - records: tuple[StoredEvent, ...] = tuple(self._stored_event(row) for row in rows[:limit]) - return EventPage( - records=records, - next_cursor=records[-1].cursor if records else None, - has_more=len(rows) > limit, - ) - - def get_run(self, run_id: str) -> RunRecord | None: - self._ensure_open() - rows: list[Mapping[str, object]] = self._read_rows( - query="SELECT * FROM sqlbuild_run_projection WHERE run_id = %s", values=[run_id] - ) - return None if not rows else self._run_record(rows[0]) - - def get_runs( - self, - *, - run_filter: RunFilter, - after_cursor: str | None = None, - limit: int = DEFAULT_PAGE_LIMIT, - ) -> RunPage: - self._ensure_open() - validate_page_limit(limit) - after_key: tuple[datetime, str] | None = self._run_cursor_key(after_cursor) - conditions: list[str] = [] - values: list[object] = [] - if after_key is not None: - conditions.append("(created_at, run_id) > (%s, %s)") - values.extend(after_key) - conditions, values = self._add_run_filters( - conditions=conditions, values=values, run_filter=run_filter - ) - where: str = " WHERE " + " AND ".join(conditions) if conditions else "" - values.append(limit + 1) - rows: list[Mapping[str, object]] = self._read_rows( - query=( - f"SELECT * FROM sqlbuild_run_projection{where} ORDER BY created_at, run_id LIMIT %s" - ), - values=values, - ) - records: tuple[RunRecord, ...] = tuple(self._run_record(row) for row in rows[:limit]) - return RunPage( - records=records, - next_cursor=self._run_cursor(records[-1]) if records else None, - has_more=len(rows) > limit, - ) - - def project(self, stored_events: Iterable[StoredEvent]) -> tuple[RunRecord, ...]: - self._ensure_open() - pending: tuple[StoredEvent, ...] = tuple(stored_events) - self._project_calls += 1 - - def publish(cursor: Any) -> tuple[RunRecord, ...]: - self._lock_history(cursor=cursor) - projected: tuple[RunRecord, ...] = project_runs( - stored_events=pending, current_runs=self._read_all_runs(cursor=cursor) - ) - self._replace_projection(cursor=cursor, projected=projected) - return projected - - return self._transactional( - action=publish, failure_message="PostgreSQL run projection publication failed" - ) - - def rebuild_from_events(self, stored_events: Iterable[StoredEvent]) -> tuple[RunRecord, ...]: - self._ensure_open() - pending: tuple[StoredEvent, ...] = tuple(stored_events) - - def rebuild(cursor: Any) -> tuple[RunRecord, ...]: - self._lock_history(cursor=cursor) - projected: tuple[RunRecord, ...] = project_runs(stored_events=pending) - self._replace_projection(cursor=cursor, projected=projected) - return projected - - return self._transactional( - action=rebuild, failure_message="PostgreSQL run projection rebuild failed" - ) - - def append_and_project( - self, events: Iterable[CanonicalLifecycleEvent] - ) -> tuple[StoredEvent, ...]: - """Atomically append immutable facts and publish their run projection.""" - - self._ensure_open() - pending: tuple[CanonicalLifecycleEvent, ...] = tuple(events) - - def append_project(cursor: Any) -> tuple[StoredEvent, ...]: - self._lock_history(cursor=cursor) - stored: tuple[StoredEvent, ...] = self._append_events(cursor=cursor, events=pending) - projected: tuple[RunRecord, ...] = project_runs( - stored_events=stored, current_runs=self._read_all_runs(cursor=cursor) - ) - self._replace_projection(cursor=cursor, projected=projected) - return stored - - return self._transactional( - action=append_project, - failure_message="PostgreSQL atomic append and projection failed", - ) - - def reconcile(self) -> tuple[RunRecord, ...]: - """Rebuild the disposable run projection from all durable event facts.""" - - self._ensure_open() - - def reconcile_projection(cursor: Any) -> tuple[RunRecord, ...]: - self._lock_history(cursor=cursor) - projected: tuple[RunRecord, ...] = project_runs( - stored_events=self._read_all_events(cursor=cursor) - ) - self._replace_projection(cursor=cursor, projected=projected) - return projected - - return self._transactional( - action=reconcile_projection, failure_message="PostgreSQL reconciliation failed" - ) - - def check_health(self) -> bool: - """Verify connectivity and the supported schema revision.""" - - self._ensure_open() - try: - with self._connection.cursor() as cursor: - cursor.execute("SELECT 1") - row: Mapping[str, object] | None = cursor.fetchone() - return row is not None and self.get_schema_version() == _SCHEMA_VERSION - except self._psycopg.Error: - return False - - def get_schema_version(self) -> int: - self._ensure_open() - return self._schema_version() - - def upgrade_schema(self, *, target_version: int | None = None) -> int: - self._ensure_open() - target: int = _SCHEMA_VERSION if target_version is None else target_version - if target != _SCHEMA_VERSION: - raise UnsupportedSchemaVersionError( - f"unsupported PostgreSQL execution history schema version {target}" - ) - self._migrate() - return target - - def close(self) -> None: - if self._closed: - return - self._connection.close() - self._closed = True - - def dispose(self) -> None: - self.close() - - def __enter__(self) -> Self: - self._ensure_open() - return self - - def __exit__( - self, - exception_type: type[BaseException] | None, - exception: BaseException | None, - traceback: TracebackType | None, - ) -> bool | None: - self.close() - return None - - @staticmethod - def _load_driver() -> tuple[Any, Any]: - try: - import psycopg - from psycopg.rows import dict_row - except ImportError: - raise ExecutionHistoryStorageError( - "PostgreSQL execution history requires psycopg 3; install sqlbuild[postgres] " - "or 'psycopg[binary]>=3.2'" - ) from None - return psycopg, dict_row - - @staticmethod - def _validate_configuration( - *, - dsn: str, - connect_timeout_seconds: int, - statement_timeout_ms: int, - max_transaction_retries: int, - ) -> None: - if not isinstance(dsn, str) or not dsn.strip(): - raise ExecutionHistoryStorageError( - "PostgreSQL execution history requires an explicit secret-resolved DSN" - ) - for name, value in ( - ("connect_timeout_seconds", connect_timeout_seconds), - ("statement_timeout_ms", statement_timeout_ms), - ): - if isinstance(value, bool) or not isinstance(value, int) or value < 1: - raise ExecutionHistoryStorageError(f"{name} must be a positive integer") - if connect_timeout_seconds > _MAX_CONNECT_TIMEOUT_SECONDS: - raise ExecutionHistoryStorageError( - f"connect_timeout_seconds must not exceed {_MAX_CONNECT_TIMEOUT_SECONDS}" - ) - if statement_timeout_ms > _MAX_STATEMENT_TIMEOUT_MS: - raise ExecutionHistoryStorageError( - f"statement_timeout_ms must not exceed {_MAX_STATEMENT_TIMEOUT_MS}" - ) - if ( - isinstance(max_transaction_retries, bool) - or not isinstance(max_transaction_retries, int) - or max_transaction_retries < 0 - or max_transaction_retries > _MAX_TRANSACTION_RETRY_LIMIT - ): - raise ExecutionHistoryStorageError( - "max_transaction_retries must be a non-negative integer no greater than " - f"{_MAX_TRANSACTION_RETRY_LIMIT}" - ) - - @staticmethod - def _close_failed_connection(*, connection: Any | None) -> None: - if connection is None: - return - try: - connection.close() - except Exception: - pass - - def _migrate(self) -> None: - def migrate(cursor: Any) -> None: - cursor.execute("SELECT pg_advisory_xact_lock(%s)", (_MIGRATION_LOCK_ID,)) - current: int = self._schema_version(cursor=cursor) - if current > _SCHEMA_VERSION: - raise UnsupportedSchemaVersionError( - f"PostgreSQL execution history schema {current} is newer than supported " - f"{_SCHEMA_VERSION}" - ) - if current < 0: - raise UnsupportedSchemaVersionError( - f"unsupported PostgreSQL execution history schema version {current}" - ) - if current == _SCHEMA_VERSION: - return - for statement in _CREATE_SCHEMA: - cursor.execute(statement) - cursor.execute( - """INSERT INTO sqlbuild_storage_migrations - (singleton, schema_version, storage_namespace) - VALUES (TRUE, %s, %s) - ON CONFLICT (singleton) DO UPDATE SET - schema_version = EXCLUDED.schema_version, - migrated_at = CURRENT_TIMESTAMP""", - (_SCHEMA_VERSION, uuid4().hex), - ) - - _ = self._transactional( - action=migrate, failure_message="PostgreSQL execution history schema migration failed" - ) - - def _schema_version(self, *, cursor: Any | None = None) -> int: - owns_cursor: bool = cursor is None - active: Any = cursor - if active is None: - active = self._connection.cursor() - try: - active.execute("SELECT to_regclass('sqlbuild_storage_migrations') AS table_name") - exists: Mapping[str, object] | None = active.fetchone() - if exists is None or exists["table_name"] is None: - return 0 - active.execute( - "SELECT schema_version FROM sqlbuild_storage_migrations WHERE singleton = TRUE" - ) - row: Mapping[str, object] | None = active.fetchone() - if row is None or isinstance(row["schema_version"], bool): - raise ExecutionHistoryStorageError( - "PostgreSQL execution history schema metadata is invalid" - ) - value: object = row["schema_version"] - if not isinstance(value, int): - raise ExecutionHistoryStorageError( - "PostgreSQL execution history schema metadata is invalid" - ) - return value - except ExecutionHistoryStorageError: - raise - except Exception: - raise ExecutionHistoryStorageError( - "PostgreSQL execution history schema metadata is invalid" - ) from None - finally: - if owns_cursor: - active.close() - - def _namespace(self, *, cursor: Any | None = None) -> str: - owns_cursor: bool = cursor is None - active: Any = cursor - if active is None: - active = self._connection.cursor() - try: - active.execute( - "SELECT storage_namespace FROM sqlbuild_storage_migrations WHERE singleton = TRUE" - ) - row: Mapping[str, object] | None = active.fetchone() - if row is None or not isinstance(row["storage_namespace"], str): - raise ExecutionHistoryStorageError( - "PostgreSQL execution history cursor namespace is missing" - ) - return row["storage_namespace"] - finally: - if owns_cursor: - active.close() - - def _transactional[Result]( - self, *, action: Callable[[Any], Result], failure_message: str - ) -> Result: - attempts = 0 - while True: - try: - with self._connection.transaction(): - with self._connection.cursor() as cursor: - cursor.execute( - "SELECT set_config('statement_timeout', %s, true)", - (str(self._statement_timeout_ms),), - ) - return action(cursor) - except (IntegrityConflictError, UnsupportedSchemaVersionError): - raise - except ExecutionHistoryStorageError: - raise - except self._psycopg.Error as error: - can_retry: bool = ( - error.sqlstate in _RETRYABLE_SQLSTATES - and attempts < self._max_transaction_retries - ) - if can_retry: - attempts += 1 - continue - raise ExecutionHistoryStorageError(failure_message) from None - - def _append_events( - self, *, cursor: Any, events: tuple[CanonicalLifecycleEvent, ...] - ) -> tuple[StoredEvent, ...]: - prepared: list[tuple[CanonicalLifecycleEvent, str, str, str]] = [] - batch_content: dict[str, str] = {} - for event in events: - event_id: str = canonical_event_id(event) - content: str = canonical_event_content(event) - digest: str = hashlib.sha256(content.encode()).hexdigest() - previous: str | None = batch_content.get(event_id) - if previous is not None and previous != digest: - raise IntegrityConflictError( - f"event_id {event_id!r} has different canonical content within batch" - ) - batch_content[event_id] = digest - prepared.append((event, event_id, content, digest)) - stored: list[StoredEvent] = [] - for event, event_id, content, digest in prepared: - envelope: dict[str, object | None] = self._event_envelope(event) - cursor.execute( - """INSERT INTO sqlbuild_lifecycle_event_log ( - event_id, schema_version, producer, producer_version, event_type, occurred_at, - invocation_id, run_id, resource_id, resource_attempt_id, operation_id, - statement_id, payload_json, content_digest - ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) - ON CONFLICT (event_id) DO NOTHING RETURNING *""", - ( - event_id, - envelope["schema_version"], - envelope["producer"], - envelope["producer_version"], - envelope["event_type"], - envelope["occurred_at"], - envelope["invocation_id"], - envelope["run_id"], - envelope["resource_id"], - envelope["resource_attempt_id"], - envelope["operation_id"], - envelope["statement_id"], - content, - digest, - ), - ) - row: Mapping[str, object] | None = cursor.fetchone() - if row is None: - cursor.execute( - "SELECT * FROM sqlbuild_lifecycle_event_log WHERE event_id = %s", - (event_id,), - ) - row = cursor.fetchone() - if row is None: - raise ExecutionHistoryStorageError( - "PostgreSQL event append did not produce a durable fact" - ) - if row["content_digest"] != digest or row["payload_json"] != content: - raise IntegrityConflictError( - f"event_id {event_id!r} has different canonical content" - ) - stored.append(self._stored_event(row)) - return tuple(stored) - - def _event_envelope(self, event: CanonicalLifecycleEvent) -> dict[str, object | None]: - if isinstance(event, LifecycleEvent): - return { - "schema_version": event.schema_version, - "producer": event.producer, - "producer_version": event.producer_version, - "event_type": event.event_type, - "occurred_at": event.occurred_at, - "invocation_id": event.invocation_id, - "run_id": event.run_id, - "resource_id": event.resource_id, - "resource_attempt_id": event.resource_attempt_id, - "operation_id": event.operation_id, - "statement_id": event.statement_id, - } - raw: Mapping[str, object] = event.raw - return { - "schema_version": self._typed_int(raw=raw, field_name="schema_version", default=1), - "producer": self._typed_text(raw=raw, field_name="producer"), - "producer_version": self._typed_text(raw=raw, field_name="producer_version"), - "event_type": self._typed_text(raw=raw, field_name="event_type"), - "occurred_at": self._opaque_timestamp(event), - "invocation_id": self._typed_text(raw=raw, field_name="invocation_id"), - "run_id": self._typed_text(raw=raw, field_name="run_id"), - "resource_id": self._typed_text(raw=raw, field_name="resource_id"), - "resource_attempt_id": self._typed_text(raw=raw, field_name="resource_attempt_id"), - "operation_id": self._typed_text(raw=raw, field_name="operation_id"), - "statement_id": self._typed_text(raw=raw, field_name="statement_id"), - } - - @staticmethod - def _typed_text(*, raw: Mapping[str, object], field_name: str) -> str | None: - value: object = raw.get(field_name) - return value if isinstance(value, str) and value else None - - @staticmethod - def _typed_int(*, raw: Mapping[str, object], field_name: str, default: int) -> int: - value: object = raw.get(field_name) - return value if isinstance(value, int) and not isinstance(value, bool) else default - - @staticmethod - def _opaque_timestamp(event: OpaqueLifecycleEvent) -> datetime | None: - value: object = event.raw.get("occurred_at") - if not isinstance(value, str): - return None - try: - parsed: datetime = datetime.fromisoformat(value.replace("Z", "+00:00")) - except ValueError: - return None - offset: timedelta | None = parsed.utcoffset() - if offset is None or offset.total_seconds() != 0: - return None - return parsed.astimezone(UTC) - - def _stored_event(self, row: Mapping[str, object]) -> StoredEvent: - raw_storage_id: object = row["storage_id"] - if not isinstance(raw_storage_id, int) or isinstance(raw_storage_id, bool): - raise ExecutionHistoryStorageError("PostgreSQL event storage position is invalid") - storage_id: int = raw_storage_id - received_at: object = row["received_at"] - if not isinstance(received_at, datetime): - raise ExecutionHistoryStorageError("PostgreSQL event timestamp is invalid") - payload_json: object = row["payload_json"] - if not isinstance(payload_json, str): - raise ExecutionHistoryStorageError("PostgreSQL canonical event content is invalid") - return StoredEvent( - storage_order=storage_id, - cursor=self._event_cursor(storage_id), - received_at=received_at.astimezone(UTC), - event=lifecycle_event_from_json(payload_json), - ) - - def _event_cursor(self, storage_id: int) -> str: - return f"{_EVENT_CURSOR_PREFIX}{self._namespace()}:{storage_id}" - - def _event_cursor_position(self, cursor: str | None) -> int: - if cursor is None: - return 0 - try: - prefix, namespace, raw_position = cursor.rsplit(":", 2) - position: int = int(raw_position) - except (AttributeError, ValueError): - raise InvalidCursorError( - "event cursor is malformed or belongs to another backend" - ) from None - if prefix + ":" != _EVENT_CURSOR_PREFIX or namespace != self._namespace() or position < 1: - raise InvalidCursorError("event cursor is not valid for this storage") - rows: list[Mapping[str, object]] = self._read_rows( - query="SELECT 1 AS present FROM sqlbuild_lifecycle_event_log WHERE storage_id = %s", - values=[position], - ) - if not rows: - raise InvalidCursorError("event cursor does not identify a durable event") - return position - - @staticmethod - def _add_event_filters( - *, conditions: list[str], values: list[object], event_filter: EventFilter - ) -> tuple[list[str], list[object]]: - filtered_conditions: list[str] = [*conditions] - filtered_values: list[object] = [*values] - for column, value in ( - ("invocation_id", event_filter.invocation_id), - ("run_id", event_filter.run_id), - ("producer", event_filter.producer), - ): - if value is not None: - filtered_conditions.append(f"{column} = %s") - filtered_values.append(value) - if event_filter.event_types: - filtered_conditions.append("event_type = ANY(%s)") - filtered_values.append(list(event_filter.event_types)) - if event_filter.family is not None: - filtered_conditions.append("event_type LIKE %s") - filtered_values.append(f"{event_filter.family.value}_%") - if event_filter.occurred_at_start is not None: - filtered_conditions.append("occurred_at >= %s") - filtered_values.append(event_filter.occurred_at_start) - if event_filter.occurred_at_end is not None: - filtered_conditions.append("occurred_at <= %s") - filtered_values.append(event_filter.occurred_at_end) - return filtered_conditions, filtered_values - - def _read_all_events(self, *, cursor: Any) -> tuple[StoredEvent, ...]: - cursor.execute("SELECT * FROM sqlbuild_lifecycle_event_log ORDER BY storage_id") - rows: list[Mapping[str, object]] = cursor.fetchall() - return tuple(self._stored_event(row) for row in rows) - - def _read_all_runs(self, *, cursor: Any) -> tuple[RunRecord, ...]: - cursor.execute("SELECT * FROM sqlbuild_run_projection ORDER BY created_at, run_id") - rows: list[Mapping[str, object]] = cursor.fetchall() - return tuple(self._run_record(row) for row in rows) - - def _replace_projection(self, *, cursor: Any, projected: tuple[RunRecord, ...]) -> None: - cursor.execute("DELETE FROM sqlbuild_run_projection") - for run in projected: - cursor.execute( - """INSERT INTO sqlbuild_run_projection ( - run_id, invocation_id, created_at, status, is_complete, last_event_cursor, - last_event_storage_id, command, target, environment, started_at, ended_at, - projection_schema_version - ) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)""", - self._run_values(run), - ) - - @staticmethod - def _run_values(run: RunRecord) -> tuple[object, ...]: - return ( - run.run_id, - run.invocation_id, - run.created_at, - run.status.value, - run.is_complete, - run.last_event_cursor, - run.last_storage_order, - run.command, - run.target, - run.environment, - run.started_at, - run.ended_at, - CURRENT_RUN_STORAGE_SCHEMA_VERSION, - ) - - @staticmethod - def _run_record(row: Mapping[str, object]) -> RunRecord: - created_at: object = row["created_at"] - started_at: object = row["started_at"] - ended_at: object = row["ended_at"] - if not isinstance(created_at, datetime): - raise ExecutionHistoryStorageError("PostgreSQL run timestamp is invalid") - last_storage_order: object = row["last_event_storage_id"] - if not isinstance(last_storage_order, int) or isinstance(last_storage_order, bool): - raise ExecutionHistoryStorageError("PostgreSQL run storage position is invalid") - return RunRecord( - run_id=str(row["run_id"]), - invocation_id=str(row["invocation_id"]), - created_at=created_at.astimezone(UTC), - status=RunStatus(str(row["status"])), - is_complete=bool(row["is_complete"]), - last_event_cursor=str(row["last_event_cursor"]), - last_storage_order=last_storage_order, - command=None if row["command"] is None else str(row["command"]), - target=None if row["target"] is None else str(row["target"]), - environment=None if row["environment"] is None else str(row["environment"]), - started_at=started_at.astimezone(UTC) if isinstance(started_at, datetime) else None, - ended_at=ended_at.astimezone(UTC) if isinstance(ended_at, datetime) else None, - ) - - def _run_cursor(self, run: RunRecord) -> str: - payload: str = json.dumps( - [self._namespace(), run.created_at.isoformat(), run.run_id], - ensure_ascii=True, - separators=(",", ":"), - ) - return _RUN_CURSOR_PREFIX + base64.urlsafe_b64encode(payload.encode()).decode() - - def _run_cursor_key(self, cursor: str | None) -> tuple[datetime, str] | None: - if cursor is None: - return None - if not cursor.startswith(_RUN_CURSOR_PREFIX): - raise InvalidCursorError("run cursor is malformed or belongs to another backend") - try: - decoded: object = json.loads( - base64.urlsafe_b64decode(cursor.removeprefix(_RUN_CURSOR_PREFIX)).decode() - ) - except (ValueError, TypeError, UnicodeError, json.JSONDecodeError, binascii.Error): - raise InvalidCursorError( - "run cursor is malformed or belongs to another backend" - ) from None - if not isinstance(decoded, list) or len(decoded) != _RUN_CURSOR_PART_COUNT: - raise InvalidCursorError("run cursor is malformed or belongs to another backend") - namespace, created_raw, run_id = decoded - if namespace != self._namespace() or not isinstance(created_raw, str): - raise InvalidCursorError("run cursor is malformed or belongs to another backend") - if not isinstance(run_id, str) or not run_id: - raise InvalidCursorError("run cursor is malformed or belongs to another backend") - try: - created_at: datetime = datetime.fromisoformat(created_raw) - except ValueError: - raise InvalidCursorError("run cursor contains an invalid timestamp") from None - if created_at.utcoffset() is None: - raise InvalidCursorError("run cursor contains an invalid timestamp") - normalized: datetime = created_at.astimezone(UTC) - rows: list[Mapping[str, object]] = self._read_rows( - query=( - "SELECT 1 AS present FROM sqlbuild_run_projection " - "WHERE created_at = %s AND run_id = %s" - ), - values=[normalized, run_id], - ) - if not rows: - raise InvalidCursorError("run cursor does not identify a run in this storage") - return normalized, run_id - - @staticmethod - def _add_run_filters( - *, conditions: list[str], values: list[object], run_filter: RunFilter - ) -> tuple[list[str], list[object]]: - filtered_conditions: list[str] = [*conditions] - filtered_values: list[object] = [*values] - if run_filter.invocation_id is not None: - filtered_conditions.append("invocation_id = %s") - filtered_values.append(run_filter.invocation_id) - if run_filter.statuses: - filtered_conditions.append("status = ANY(%s)") - filtered_values.append([status.value for status in run_filter.statuses]) - if run_filter.created_at_start is not None: - filtered_conditions.append("created_at >= %s") - filtered_values.append(run_filter.created_at_start) - if run_filter.created_at_end is not None: - filtered_conditions.append("created_at <= %s") - filtered_values.append(run_filter.created_at_end) - return filtered_conditions, filtered_values - - def _read_rows(self, *, query: str, values: list[object]) -> list[Mapping[str, object]]: - try: - with self._connection.cursor() as cursor: - cursor.execute(query, values) - return cursor.fetchall() - except self._psycopg.Error: - raise ExecutionHistoryStorageError("PostgreSQL execution history read failed") from None - - @staticmethod - def _lock_history(*, cursor: Any) -> None: - cursor.execute("SELECT pg_advisory_xact_lock(%s)", (_HISTORY_LOCK_ID,)) - - def _ensure_open(self) -> None: - if self._closed: - raise ExecutionHistoryStorageError("PostgreSQL execution history storage is closed") diff --git a/src/sqlbuild/runtime/execution_history/classes/sqlite_execution_history.py b/src/sqlbuild/runtime/execution_history/classes/sqlite_execution_history.py deleted file mode 100644 index c1d464addb..0000000000 --- a/src/sqlbuild/runtime/execution_history/classes/sqlite_execution_history.py +++ /dev/null @@ -1,795 +0,0 @@ -"""Project-local SQLite execution history storage.""" - -from __future__ import annotations - -import base64 -import binascii -import hashlib -import json -import os -import sqlite3 -import threading -from collections.abc import Iterable, Iterator, Mapping -from contextlib import contextmanager -from datetime import UTC, datetime, timedelta -from pathlib import Path -from types import TracebackType -from typing import Self -from uuid import uuid4 - -from sqlbuild.runtime.execution_history.constants import ( - CURRENT_RUN_STORAGE_SCHEMA_VERSION, - DEFAULT_PAGE_LIMIT, -) -from sqlbuild.runtime.execution_history.exceptions import ( - ExecutionHistoryStorageError, - IntegrityConflictError, - InvalidCursorError, - UnsupportedSchemaVersionError, -) -from sqlbuild.runtime.execution_history.main.canonical_event_content import canonical_event_content -from sqlbuild.runtime.execution_history.main.canonical_event_id import canonical_event_id -from sqlbuild.runtime.execution_history.main.project_runs import project_runs -from sqlbuild.runtime.execution_history.main.validate_page_limit import validate_page_limit -from sqlbuild.runtime.execution_history.models import ( - EventFilter, - EventPage, - RunFilter, - RunPage, - RunRecord, - StoredEvent, -) -from sqlbuild.runtime.execution_history.types import CanonicalLifecycleEvent, RunStatus -from sqlbuild.runtime.observability.main.lifecycle_event_from_json import lifecycle_event_from_json -from sqlbuild.runtime.observability.models import LifecycleEvent, OpaqueLifecycleEvent - -_BUSY_TIMEOUT_MS: int = 5_000 -_CURSOR_PREFIX: str = "sqlite-event:" -_RUN_CURSOR_PREFIX: str = "sqlite-run:" -_SCHEMA_VERSION: int = 1 -_MEMORY_PATH: str = ":memory:" -_HEALTHY_CHECK_RESULT: str = "ok" -_RUN_CURSOR_PART_COUNT: int = 3 - -_CREATE_SCHEMA: str = """ -CREATE TABLE IF NOT EXISTS execution_history_metadata ( - key TEXT PRIMARY KEY, - value TEXT NOT NULL -); -CREATE TABLE IF NOT EXISTS lifecycle_event_log ( - storage_id INTEGER PRIMARY KEY AUTOINCREMENT, - event_id TEXT NOT NULL UNIQUE, - schema_version INTEGER NOT NULL, - producer TEXT, - event_type TEXT, - occurred_at TEXT, - received_at TEXT NOT NULL, - invocation_id TEXT, - run_id TEXT, - resource_id TEXT, - resource_attempt_id TEXT, - operation_id TEXT, - statement_id TEXT, - payload_json TEXT NOT NULL, - content_digest TEXT NOT NULL -); -CREATE INDEX IF NOT EXISTS lifecycle_event_log_run_storage - ON lifecycle_event_log (run_id, storage_id); -CREATE INDEX IF NOT EXISTS lifecycle_event_log_invocation_storage - ON lifecycle_event_log (invocation_id, storage_id); -CREATE INDEX IF NOT EXISTS lifecycle_event_log_type_storage - ON lifecycle_event_log (event_type, storage_id); -CREATE TABLE IF NOT EXISTS run_projection ( - run_id TEXT PRIMARY KEY, - invocation_id TEXT NOT NULL, - created_at TEXT NOT NULL, - status TEXT NOT NULL, - is_complete INTEGER NOT NULL, - last_event_cursor TEXT NOT NULL, - last_storage_order INTEGER NOT NULL, - command TEXT, - target TEXT, - environment TEXT, - started_at TEXT, - ended_at TEXT, - projection_schema_version INTEGER NOT NULL -); -CREATE INDEX IF NOT EXISTS run_projection_created ON run_projection (created_at, run_id); -CREATE INDEX IF NOT EXISTS run_projection_invocation_created - ON run_projection (invocation_id, created_at, run_id); -""" - - -class SQLiteExecutionHistory: - """SQLite implementation of event-log and run-projection storage contracts.""" - - def __init__( - self, - project_dir: Path | str | None = None, - *, - path: Path | str | None = None, - busy_timeout_ms: int = _BUSY_TIMEOUT_MS, - ) -> None: - if project_dir is not None and path is not None: - raise ExecutionHistoryStorageError("provide project_dir or path, not both") - if isinstance(busy_timeout_ms, bool) or not isinstance(busy_timeout_ms, int): - raise ExecutionHistoryStorageError("busy_timeout_ms must be a positive integer") - if busy_timeout_ms < 1: - raise ExecutionHistoryStorageError("busy_timeout_ms must be a positive integer") - selected: Path | str = ( - path if path is not None else Path.cwd() if project_dir is None else project_dir - ) - self._path: str = ( - selected - if isinstance(selected, str) and selected == _MEMORY_PATH - else str(Path(selected) / ".sqlbuild" / "history.sqlite3") - if path is None - else str(Path(selected)) - ) - self._closed = False - self._project_calls = 0 - self._lock: threading.RLock = threading.RLock() - connection: sqlite3.Connection | None = None - try: - self._prepare_parent() - connection = sqlite3.connect( - self._path, - timeout=busy_timeout_ms / 1000, - isolation_level=None, - check_same_thread=False, - ) - self._connection: sqlite3.Connection = connection - with self._lock: - self._connection.row_factory = sqlite3.Row - _ = self._connection.execute(f"PRAGMA busy_timeout = {busy_timeout_ms}") - _ = self._connection.execute("PRAGMA foreign_keys = ON") - if self._path != _MEMORY_PATH: - _ = self._connection.execute("PRAGMA journal_mode = WAL").fetchone() - self._migrate() - self._restrict_file_permissions() - _ = self.reconcile() - except ExecutionHistoryStorageError: - self._close_failed_connection(connection=connection) - self._closed = True - raise - except (OSError, sqlite3.Error) as error: - self._close_failed_connection(connection=connection) - self._closed = True - raise ExecutionHistoryStorageError( - f"cannot open SQLite history at {self._path}" - ) from error - - @property - def path(self) -> Path | None: - """Return the database path, or None for an in-memory database.""" - - return None if self._path == _MEMORY_PATH else Path(self._path) - - @property - def project_calls(self) -> int: - """Return the number of incremental projection calls.""" - - with self._lock: - return self._project_calls - - def append_event(self, event: CanonicalLifecycleEvent) -> StoredEvent: - with self._lock: - self._ensure_open() - return self.append_events((event,))[0] - - def append_events(self, events: Iterable[CanonicalLifecycleEvent]) -> tuple[StoredEvent, ...]: - with self._lock: - self._ensure_open() - pending: tuple[CanonicalLifecycleEvent, ...] = tuple(events) - try: - with self._transaction(): - return self._append_events(pending) - except IntegrityConflictError: - raise - except sqlite3.Error as error: - raise ExecutionHistoryStorageError("SQLite event append failed") from error - - def get_events( - self, - *, - event_filter: EventFilter, - after_cursor: str | None = None, - limit: int = DEFAULT_PAGE_LIMIT, - ) -> EventPage: - with self._lock: - self._ensure_open() - validate_page_limit(limit) - after_storage_id: int = self._event_cursor_position(after_cursor) - conditions: list[str] = ["storage_id > ?"] - values: list[object] = [after_storage_id] - conditions, values = self._add_event_filters( - conditions=conditions, values=values, event_filter=event_filter - ) - sql: str = ( - "SELECT * FROM lifecycle_event_log WHERE " - + " AND ".join(conditions) - + " ORDER BY storage_id ASC LIMIT ?" - ) - values.append(limit + 1) - try: - rows: list[sqlite3.Row] = list(self._connection.execute(sql, values)) - except sqlite3.Error as error: - raise ExecutionHistoryStorageError("SQLite event read failed") from error - records: tuple[StoredEvent, ...] = tuple( - self._stored_event(row) for row in rows[:limit] - ) - return EventPage( - records=records, - next_cursor=records[-1].cursor if records else None, - has_more=len(rows) > limit, - ) - - def get_run(self, run_id: str) -> RunRecord | None: - with self._lock: - self._ensure_open() - row: sqlite3.Row | None = self._connection.execute( - "SELECT * FROM run_projection WHERE run_id = ?", (run_id,) - ).fetchone() - return None if row is None else self._run_record(row) - - def get_runs( - self, - *, - run_filter: RunFilter, - after_cursor: str | None = None, - limit: int = DEFAULT_PAGE_LIMIT, - ) -> RunPage: - with self._lock: - self._ensure_open() - validate_page_limit(limit) - after_key: tuple[str, str] | None = self._run_cursor_key(after_cursor) - conditions: list[str] = [] - values: list[object] = [] - if after_key is not None: - conditions.append("(created_at > ? OR (created_at = ? AND run_id > ?))") - values.extend((after_key[0], after_key[0], after_key[1])) - conditions, values = self._add_run_filters( - conditions=conditions, values=values, run_filter=run_filter - ) - where: str = " WHERE " + " AND ".join(conditions) if conditions else "" - values.append(limit + 1) - rows: list[sqlite3.Row] = list( - self._connection.execute( - f"SELECT * FROM run_projection{where} ORDER BY created_at, run_id LIMIT ?", - values, - ) - ) - records: tuple[RunRecord, ...] = tuple(self._run_record(row) for row in rows[:limit]) - return RunPage( - records=records, - next_cursor=self._run_cursor(records[-1]) if records else None, - has_more=len(rows) > limit, - ) - - def project(self, stored_events: Iterable[StoredEvent]) -> tuple[RunRecord, ...]: - with self._lock: - self._ensure_open() - self._project_calls += 1 - current: tuple[RunRecord, ...] = self._read_all_runs() - projected: tuple[RunRecord, ...] = project_runs( - stored_events=stored_events, current_runs=current - ) - self._publish_projection(projected) - return projected - - def rebuild_from_events(self, stored_events: Iterable[StoredEvent]) -> tuple[RunRecord, ...]: - with self._lock: - self._ensure_open() - projected: tuple[RunRecord, ...] = project_runs(stored_events=stored_events) - self._publish_projection(projected) - return projected - - def append_and_project( - self, events: Iterable[CanonicalLifecycleEvent] - ) -> tuple[StoredEvent, ...]: - """Atomically append immutable facts and publish their run projection.""" - - with self._lock: - self._ensure_open() - pending: tuple[CanonicalLifecycleEvent, ...] = tuple(events) - try: - with self._transaction(): - stored: tuple[StoredEvent, ...] = self._append_events(pending) - projected: tuple[RunRecord, ...] = project_runs( - stored_events=stored, current_runs=self._read_all_runs() - ) - self._replace_projection(projected) - return stored - except IntegrityConflictError: - raise - except sqlite3.Error as error: - raise ExecutionHistoryStorageError( - "SQLite atomic append and projection failed" - ) from error - - def reconcile(self) -> tuple[RunRecord, ...]: - """Rebuild the disposable run projection from all durable event facts.""" - - with self._lock: - self._ensure_open() - try: - with self._transaction(): - records: tuple[StoredEvent, ...] = self._read_all_events() - projected: tuple[RunRecord, ...] = project_runs(stored_events=records) - self._replace_projection(projected) - return projected - except sqlite3.Error as error: - raise ExecutionHistoryStorageError("SQLite reconciliation failed") from error - - def check_health(self) -> bool: - """Verify connectivity, integrity, and the supported schema revision.""" - - with self._lock: - self._ensure_open() - result: sqlite3.Row | None = self._connection.execute("PRAGMA quick_check").fetchone() - return ( - result is not None - and result[0] == _HEALTHY_CHECK_RESULT - and self.get_schema_version() == _SCHEMA_VERSION - ) - - def get_schema_version(self) -> int: - with self._lock: - self._ensure_open() - return self._metadata_version() - - def upgrade_schema(self, *, target_version: int | None = None) -> int: - with self._lock: - self._ensure_open() - target: int = _SCHEMA_VERSION if target_version is None else target_version - if target != _SCHEMA_VERSION: - raise UnsupportedSchemaVersionError(f"unsupported SQLite schema version {target}") - self._migrate() - return target - - def close(self) -> None: - with self._lock: - if self._closed: - return - try: - self._connection.close() - finally: - self._closed = True - - def dispose(self) -> None: - self.close() - - def __enter__(self) -> Self: - with self._lock: - self._ensure_open() - return self - - def __exit__( - self, - exception_type: type[BaseException] | None, - exception: BaseException | None, - traceback: TracebackType | None, - ) -> bool | None: - self.close() - return None - - def _prepare_parent(self) -> None: - if self._path == _MEMORY_PATH: - return - parent: Path = Path(self._path).parent - parent.mkdir(mode=0o700, parents=True, exist_ok=True) - - def _restrict_file_permissions(self) -> None: - if self._path != _MEMORY_PATH: - os.chmod(self._path, 0o600) - - @staticmethod - def _close_failed_connection(*, connection: sqlite3.Connection | None) -> None: - if connection is None: - return - try: - connection.close() - except sqlite3.Error: - pass - - def _migrate(self) -> None: - current: int = self._metadata_version() - if current > _SCHEMA_VERSION: - raise UnsupportedSchemaVersionError( - f"SQLite history schema {current} is newer than supported {_SCHEMA_VERSION}" - ) - if current < 0: - raise UnsupportedSchemaVersionError(f"unsupported SQLite schema version {current}") - if current == _SCHEMA_VERSION: - return - try: - namespace: str = uuid4().hex - with self._transaction(): - for statement in _CREATE_SCHEMA.split(";"): - if statement.strip(): - _ = self._connection.execute(statement) - _ = self._connection.execute( - "INSERT OR IGNORE INTO execution_history_metadata(key, value) VALUES (?, ?)", - ("storage_namespace", namespace), - ) - _ = self._connection.execute( - "INSERT OR REPLACE INTO execution_history_metadata(key, value) VALUES (?, ?)", - ("schema_version", str(_SCHEMA_VERSION)), - ) - except sqlite3.Error as error: - raise ExecutionHistoryStorageError("SQLite history schema migration failed") from error - - def _metadata_version(self) -> int: - exists: sqlite3.Row | None = self._connection.execute( - "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?", - ("execution_history_metadata",), - ).fetchone() - if exists is None: - return 0 - row: sqlite3.Row | None = self._connection.execute( - "SELECT value FROM execution_history_metadata WHERE key = ?", ("schema_version",) - ).fetchone() - if row is None: - return 0 - try: - return int(row[0]) - except (TypeError, ValueError) as error: - raise ExecutionHistoryStorageError( - "SQLite history schema metadata is invalid" - ) from error - - def _namespace(self) -> str: - row: sqlite3.Row | None = self._connection.execute( - "SELECT value FROM execution_history_metadata WHERE key = ?", ("storage_namespace",) - ).fetchone() - if row is None: - raise ExecutionHistoryStorageError("SQLite history cursor namespace is missing") - return str(row[0]) - - def _append_events( - self, events: tuple[CanonicalLifecycleEvent, ...] - ) -> tuple[StoredEvent, ...]: - prepared: list[tuple[CanonicalLifecycleEvent, str, str, str]] = [] - batch_content: dict[str, str] = {} - for event in events: - event_id: str = canonical_event_id(event) - content: str = canonical_event_content(event) - digest: str = hashlib.sha256(content.encode()).hexdigest() - previous: str | None = batch_content.get(event_id) - if previous is not None and previous != digest: - raise IntegrityConflictError( - f"event_id {event_id!r} has different canonical content within batch" - ) - batch_content[event_id] = digest - prepared.append((event, event_id, content, digest)) - stored: list[StoredEvent] = [] - for event, event_id, content, digest in prepared: - existing: sqlite3.Row | None = self._connection.execute( - "SELECT * FROM lifecycle_event_log WHERE event_id = ?", (event_id,) - ).fetchone() - if existing is not None: - if existing["content_digest"] != digest or existing["payload_json"] != content: - raise IntegrityConflictError( - f"event_id {event_id!r} has different canonical content" - ) - stored.append(self._stored_event(existing)) - continue - received_at: str = self._timestamp(datetime.now(UTC)) - envelope: dict[str, object | None] = self._event_envelope(event) - cursor: sqlite3.Cursor = self._connection.execute( - """INSERT INTO lifecycle_event_log ( - event_id, schema_version, producer, event_type, occurred_at, received_at, - invocation_id, run_id, resource_id, resource_attempt_id, operation_id, - statement_id, payload_json, content_digest - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", - ( - event_id, - envelope["schema_version"], - envelope["producer"], - envelope["event_type"], - envelope["occurred_at"], - received_at, - envelope["invocation_id"], - envelope["run_id"], - envelope["resource_id"], - envelope["resource_attempt_id"], - envelope["operation_id"], - envelope["statement_id"], - content, - digest, - ), - ) - row: sqlite3.Row = self._connection.execute( - "SELECT * FROM lifecycle_event_log WHERE storage_id = ?", (cursor.lastrowid,) - ).fetchone() - stored.append(self._stored_event(row)) - return tuple(stored) - - def _event_envelope(self, event: CanonicalLifecycleEvent) -> dict[str, object | None]: - if isinstance(event, LifecycleEvent): - return { - "schema_version": event.schema_version, - "producer": event.producer, - "event_type": event.event_type, - "occurred_at": self._timestamp(event.occurred_at), - "invocation_id": event.invocation_id, - "run_id": event.run_id, - "resource_id": event.resource_id, - "resource_attempt_id": event.resource_attempt_id, - "operation_id": event.operation_id, - "statement_id": event.statement_id, - } - raw: Mapping[str, object] = event.raw - return { - "schema_version": self._typed_int(raw=raw, field_name="schema_version", default=1), - "producer": self._typed_text(raw=raw, field_name="producer"), - "event_type": self._typed_text(raw=raw, field_name="event_type"), - "occurred_at": self._opaque_timestamp(event), - "invocation_id": self._typed_text(raw=raw, field_name="invocation_id"), - "run_id": self._typed_text(raw=raw, field_name="run_id"), - "resource_id": self._typed_text(raw=raw, field_name="resource_id"), - "resource_attempt_id": self._typed_text(raw=raw, field_name="resource_attempt_id"), - "operation_id": self._typed_text(raw=raw, field_name="operation_id"), - "statement_id": self._typed_text(raw=raw, field_name="statement_id"), - } - - @staticmethod - def _typed_text(*, raw: Mapping[str, object], field_name: str) -> str | None: - value: object = raw.get(field_name) - return value if isinstance(value, str) and value else None - - @staticmethod - def _typed_int(*, raw: Mapping[str, object], field_name: str, default: int) -> int: - value: object = raw.get(field_name) - return value if isinstance(value, int) and not isinstance(value, bool) else default - - @staticmethod - def _opaque_timestamp(event: OpaqueLifecycleEvent) -> str | None: - value: object = event.raw.get("occurred_at") - if not isinstance(value, str): - return None - try: - parsed: datetime = datetime.fromisoformat(value.replace("Z", "+00:00")) - except ValueError: - return None - offset: timedelta | None = parsed.utcoffset() - if offset is None or offset.total_seconds() != 0: - return None - return SQLiteExecutionHistory._timestamp(parsed) - - def _stored_event(self, row: sqlite3.Row) -> StoredEvent: - storage_id: int = int(row["storage_id"]) - return StoredEvent( - storage_order=storage_id, - cursor=self._event_cursor(storage_id), - received_at=self._parse_timestamp(row["received_at"]), - event=lifecycle_event_from_json(row["payload_json"]), - ) - - def _event_cursor(self, storage_id: int) -> str: - return f"{_CURSOR_PREFIX}{self._namespace()}:{storage_id}" - - def _event_cursor_position(self, cursor: str | None) -> int: - if cursor is None: - return 0 - try: - prefix, namespace, raw_position = cursor.rsplit(":", 2) - position: int = int(raw_position) - except (AttributeError, ValueError) as error: - raise InvalidCursorError( - "event cursor is malformed or belongs to another backend" - ) from error - if prefix + ":" != _CURSOR_PREFIX or namespace != self._namespace() or position < 1: - raise InvalidCursorError("event cursor is not valid for this storage") - exists: sqlite3.Row | None = self._connection.execute( - "SELECT 1 FROM lifecycle_event_log WHERE storage_id = ?", (position,) - ).fetchone() - if exists is None: - raise InvalidCursorError("event cursor does not identify a durable event") - return position - - @staticmethod - def _add_event_filters( - *, conditions: list[str], values: list[object], event_filter: EventFilter - ) -> tuple[list[str], list[object]]: - filtered_conditions: list[str] = [*conditions] - filtered_values: list[object] = [*values] - for column, value in ( - ("invocation_id", event_filter.invocation_id), - ("run_id", event_filter.run_id), - ("producer", event_filter.producer), - ): - if value is not None: - filtered_conditions.append(f"{column} = ?") - filtered_values.append(value) - if event_filter.event_types: - placeholders: str = ",".join("?" for _ in event_filter.event_types) - filtered_conditions.append(f"event_type IN ({placeholders})") - filtered_values.extend(event_filter.event_types) - if event_filter.family is not None: - filtered_conditions.append("event_type LIKE ?") - filtered_values.append(f"{event_filter.family.value}_%") - if event_filter.occurred_at_start is not None: - filtered_conditions.append("occurred_at >= ?") - filtered_values.append( - SQLiteExecutionHistory._timestamp(event_filter.occurred_at_start) - ) - if event_filter.occurred_at_end is not None: - filtered_conditions.append("occurred_at <= ?") - filtered_values.append(SQLiteExecutionHistory._timestamp(event_filter.occurred_at_end)) - return filtered_conditions, filtered_values - - def _read_all_events(self) -> tuple[StoredEvent, ...]: - rows: list[sqlite3.Row] = list( - self._connection.execute("SELECT * FROM lifecycle_event_log ORDER BY storage_id") - ) - return tuple(self._stored_event(row) for row in rows) - - def _read_all_runs(self) -> tuple[RunRecord, ...]: - rows: list[sqlite3.Row] = list( - self._connection.execute("SELECT * FROM run_projection ORDER BY created_at, run_id") - ) - return tuple(self._run_record(row) for row in rows) - - def _publish_projection(self, projected: tuple[RunRecord, ...]) -> None: - try: - with self._transaction(): - self._replace_projection(projected) - except sqlite3.Error as error: - raise ExecutionHistoryStorageError( - "SQLite run projection publication failed" - ) from error - - def _replace_projection(self, projected: tuple[RunRecord, ...]) -> None: - _ = self._connection.execute("DELETE FROM run_projection") - self._connection.executemany( - """INSERT INTO run_projection ( - run_id, invocation_id, created_at, status, is_complete, last_event_cursor, - last_storage_order, command, target, environment, started_at, ended_at, - projection_schema_version - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""", - tuple(self._run_values(run) for run in projected), - ) - - @staticmethod - def _run_values(run: RunRecord) -> tuple[object, ...]: - return ( - run.run_id, - run.invocation_id, - SQLiteExecutionHistory._timestamp(run.created_at), - run.status.value, - int(run.is_complete), - run.last_event_cursor, - run.last_storage_order, - run.command, - run.target, - run.environment, - None if run.started_at is None else SQLiteExecutionHistory._timestamp(run.started_at), - None if run.ended_at is None else SQLiteExecutionHistory._timestamp(run.ended_at), - CURRENT_RUN_STORAGE_SCHEMA_VERSION, - ) - - @staticmethod - def _run_record(row: sqlite3.Row) -> RunRecord: - return RunRecord( - run_id=row["run_id"], - invocation_id=row["invocation_id"], - created_at=SQLiteExecutionHistory._parse_timestamp(row["created_at"]), - status=RunStatus(row["status"]), - is_complete=bool(row["is_complete"]), - last_event_cursor=row["last_event_cursor"], - last_storage_order=int(row["last_storage_order"]), - command=row["command"], - target=row["target"], - environment=row["environment"], - started_at=( - None - if row["started_at"] is None - else SQLiteExecutionHistory._parse_timestamp(row["started_at"]) - ), - ended_at=( - None - if row["ended_at"] is None - else SQLiteExecutionHistory._parse_timestamp(row["ended_at"]) - ), - ) - - def _run_cursor(self, run: RunRecord) -> str: - payload: str = json.dumps( - [self._namespace(), self._timestamp(run.created_at), run.run_id], - ensure_ascii=True, - separators=(",", ":"), - ) - return _RUN_CURSOR_PREFIX + base64.urlsafe_b64encode(payload.encode()).decode() - - def _run_cursor_key(self, cursor: str | None) -> tuple[str, str] | None: - if cursor is None: - return None - if not cursor.startswith(_RUN_CURSOR_PREFIX): - raise InvalidCursorError("run cursor is malformed or belongs to another backend") - try: - decoded: object = json.loads( - base64.urlsafe_b64decode(cursor.removeprefix(_RUN_CURSOR_PREFIX)).decode() - ) - except (ValueError, TypeError, UnicodeError, json.JSONDecodeError, binascii.Error) as error: - raise InvalidCursorError( - "run cursor is malformed or belongs to another backend" - ) from error - if not isinstance(decoded, list) or len(decoded) != _RUN_CURSOR_PART_COUNT: - raise InvalidCursorError("run cursor is malformed or belongs to another backend") - namespace, created_at, run_id = decoded - if namespace != self._namespace() or not isinstance(created_at, str): - raise InvalidCursorError("run cursor is malformed or belongs to another backend") - if not isinstance(run_id, str) or not run_id: - raise InvalidCursorError("run cursor is malformed or belongs to another backend") - try: - _ = self._parse_timestamp(created_at) - except (ValueError, ExecutionHistoryStorageError) as error: - raise InvalidCursorError("run cursor contains an invalid timestamp") from error - exists: sqlite3.Row | None = self._connection.execute( - "SELECT 1 FROM run_projection WHERE created_at = ? AND run_id = ?", - (created_at, run_id), - ).fetchone() - if exists is None: - raise InvalidCursorError("run cursor does not identify a run in this storage") - return created_at, run_id - - @staticmethod - def _add_run_filters( - *, conditions: list[str], values: list[object], run_filter: RunFilter - ) -> tuple[list[str], list[object]]: - filtered_conditions: list[str] = [*conditions] - filtered_values: list[object] = [*values] - if run_filter.invocation_id is not None: - filtered_conditions.append("invocation_id = ?") - filtered_values.append(run_filter.invocation_id) - if run_filter.statuses: - placeholders: str = ",".join("?" for _ in run_filter.statuses) - filtered_conditions.append(f"status IN ({placeholders})") - filtered_values.extend(status.value for status in run_filter.statuses) - if run_filter.created_at_start is not None: - filtered_conditions.append("created_at >= ?") - filtered_values.append(SQLiteExecutionHistory._timestamp(run_filter.created_at_start)) - if run_filter.created_at_end is not None: - filtered_conditions.append("created_at <= ?") - filtered_values.append(SQLiteExecutionHistory._timestamp(run_filter.created_at_end)) - return filtered_conditions, filtered_values - - @staticmethod - def _timestamp(value: datetime) -> str: - return value.astimezone(UTC).isoformat(timespec="microseconds") - - @staticmethod - def _parse_timestamp(value: str) -> datetime: - parsed: datetime = datetime.fromisoformat(value.replace("Z", "+00:00")) - if parsed.utcoffset() is None: - raise ExecutionHistoryStorageError("SQLite history timestamp is not timezone-aware") - return parsed.astimezone(UTC) - - @contextmanager - def _transaction(self) -> Iterator[None]: - with self._lock: - self._ensure_open() - _ = self._connection.execute("BEGIN IMMEDIATE") - try: - yield - except BaseException: - try: - _ = self._connection.execute("ROLLBACK") - except BaseException: - pass - raise - else: - try: - _ = self._connection.execute("COMMIT") - except sqlite3.Error: - if self._connection.in_transaction: - try: - _ = self._connection.execute("ROLLBACK") - except BaseException: - pass - raise - - def _ensure_open(self) -> None: - if self._closed: - raise ExecutionHistoryStorageError("SQLite execution history storage is closed") diff --git a/src/sqlbuild/runtime/execution_history/main/__init__.py b/src/sqlbuild/runtime/execution_history/main/__init__.py deleted file mode 100644 index c066fd8e21..0000000000 --- a/src/sqlbuild/runtime/execution_history/main/__init__.py +++ /dev/null @@ -1 +0,0 @@ -"""Execution history public entry implementations.""" diff --git a/src/sqlbuild/runtime/execution_history/main/append_and_project.py b/src/sqlbuild/runtime/execution_history/main/append_and_project.py deleted file mode 100644 index 7c95b78e92..0000000000 --- a/src/sqlbuild/runtime/execution_history/main/append_and_project.py +++ /dev/null @@ -1,23 +0,0 @@ -"""Durable append and rebuildable projection orchestration.""" - -from collections.abc import Iterable - -from sqlbuild.runtime.execution_history.models import StoredEvent -from sqlbuild.runtime.execution_history.types import ( - CanonicalLifecycleEvent, - LifecycleEventLogStorage, - RunStorage, -) - - -def append_and_project( - *, - event_log: LifecycleEventLogStorage, - run_storage: RunStorage, - events: Iterable[CanonicalLifecycleEvent], -) -> tuple[StoredEvent, ...]: - """Append facts durably before publishing their rebuildable run projection.""" - - stored_events: tuple[StoredEvent, ...] = event_log.append_events(events) - _ = run_storage.project(stored_events) - return stored_events diff --git a/src/sqlbuild/runtime/execution_history/main/canonical_event_content.py b/src/sqlbuild/runtime/execution_history/main/canonical_event_content.py deleted file mode 100644 index 4e930253fd..0000000000 --- a/src/sqlbuild/runtime/execution_history/main/canonical_event_content.py +++ /dev/null @@ -1,12 +0,0 @@ -"""Canonical lifecycle event content entrypoint.""" - -from sqlbuild.runtime.execution_history._helpers.canonical_events import ( - canonical_event_content as _canonical_event_content, -) -from sqlbuild.runtime.execution_history.types import CanonicalLifecycleEvent - - -def canonical_event_content(event: CanonicalLifecycleEvent) -> str: - """Return deterministic content for lifecycle event idempotency comparison.""" - - return _canonical_event_content(event=event) diff --git a/src/sqlbuild/runtime/execution_history/main/canonical_event_id.py b/src/sqlbuild/runtime/execution_history/main/canonical_event_id.py deleted file mode 100644 index d7c1cb1b37..0000000000 --- a/src/sqlbuild/runtime/execution_history/main/canonical_event_id.py +++ /dev/null @@ -1,12 +0,0 @@ -"""Canonical lifecycle event ID entrypoint.""" - -from sqlbuild.runtime.execution_history._helpers.canonical_events import ( - canonical_event_id as _canonical_event_id, -) -from sqlbuild.runtime.execution_history.types import CanonicalLifecycleEvent - - -def canonical_event_id(event: CanonicalLifecycleEvent) -> str: - """Return the required stable ID from a known or opaque lifecycle event.""" - - return _canonical_event_id(event=event) diff --git a/src/sqlbuild/runtime/execution_history/main/project_runs.py b/src/sqlbuild/runtime/execution_history/main/project_runs.py deleted file mode 100644 index 3b0602d217..0000000000 --- a/src/sqlbuild/runtime/execution_history/main/project_runs.py +++ /dev/null @@ -1,14 +0,0 @@ -"""Pure run projection entrypoint.""" - -from collections.abc import Iterable - -from sqlbuild.runtime.execution_history._helpers.projection import project_runs as _project_runs -from sqlbuild.runtime.execution_history.models import RunRecord, StoredEvent - - -def project_runs( - *, stored_events: Iterable[StoredEvent], current_runs: Iterable[RunRecord] = () -) -> tuple[RunRecord, ...]: - """Apply durable event facts in storage order to immutable run projections.""" - - return _project_runs(stored_events=stored_events, current_runs=current_runs) diff --git a/src/sqlbuild/runtime/execution_history/main/validate_page_limit.py b/src/sqlbuild/runtime/execution_history/main/validate_page_limit.py deleted file mode 100644 index cfb55c5db4..0000000000 --- a/src/sqlbuild/runtime/execution_history/main/validate_page_limit.py +++ /dev/null @@ -1,11 +0,0 @@ -"""Storage page limit validation entrypoint.""" - -from sqlbuild.runtime.execution_history._helpers.validation import ( - validate_page_limit as _validate_page_limit, -) - - -def validate_page_limit(limit: int) -> None: - """Validate a positive page limit no greater than the public maximum.""" - - _validate_page_limit(limit) diff --git a/src/sqlbuild/runtime/execution_history/models.py b/src/sqlbuild/runtime/execution_history/models.py index 38d56729d8..2b4ecf920f 100644 --- a/src/sqlbuild/runtime/execution_history/models.py +++ b/src/sqlbuild/runtime/execution_history/models.py @@ -5,9 +5,11 @@ from dataclasses import dataclass from datetime import datetime -from sqlbuild.runtime.execution_history.exceptions import ( - InvalidFilterError, - ProjectionConsistencyError, +from sqlbuild._native import ( + validate_history_event_filter, + validate_history_page, + validate_history_run_filter, + validate_stored_history_event, ) from sqlbuild.runtime.execution_history.types import ( CanonicalLifecycleEvent, @@ -28,15 +30,7 @@ class StoredEvent: event: CanonicalLifecycleEvent def __post_init__(self) -> None: - from sqlbuild.runtime.execution_history._helpers.validation import ( - validate_storage_timestamp, - ) - - if self.storage_order < 1: - raise ProjectionConsistencyError("storage_order must be a positive integer") - if not self.cursor: - raise ProjectionConsistencyError("event cursor must be non-empty") - validate_storage_timestamp(value=self.received_at, field_name="received_at") + validate_stored_history_event(self) @dataclass(frozen=True) @@ -52,28 +46,7 @@ class EventFilter: occurred_at_end: datetime | None = None def __post_init__(self) -> None: - from sqlbuild.runtime.execution_history._helpers.validation import ( - validate_filter_text, - validate_filter_timestamp, - ) - - for field_name in ("invocation_id", "run_id", "producer"): - validate_filter_text(value=getattr(self, field_name), field_name=field_name) - if self.event_types and self.family is not None: - raise InvalidFilterError("event_types and family are mutually exclusive") - if any(not event_type.strip() for event_type in self.event_types): - raise InvalidFilterError("event_types must contain only non-empty values") - for field_name, value in ( - ("occurred_at_start", self.occurred_at_start), - ("occurred_at_end", self.occurred_at_end), - ): - validate_filter_timestamp(value=value, field_name=field_name) - if ( - self.occurred_at_start is not None - and self.occurred_at_end is not None - and self.occurred_at_start > self.occurred_at_end - ): - raise InvalidFilterError("occurred_at_start must not be after occurred_at_end") + validate_history_event_filter(self) @dataclass(frozen=True) @@ -85,10 +58,7 @@ class EventPage: has_more: bool def __post_init__(self) -> None: - if not self.records and self.next_cursor is not None: - raise ProjectionConsistencyError("an empty event page cannot have a next cursor") - if self.records and self.next_cursor != self.records[-1].cursor: - raise ProjectionConsistencyError("event page cursor must identify its last record") + validate_history_page(page=self, is_event_page=True) @dataclass(frozen=True) @@ -119,23 +89,7 @@ class RunFilter: created_at_end: datetime | None = None def __post_init__(self) -> None: - from sqlbuild.runtime.execution_history._helpers.validation import ( - validate_filter_text, - validate_filter_timestamp, - ) - - validate_filter_text(value=self.invocation_id, field_name="invocation_id") - for field_name, value in ( - ("created_at_start", self.created_at_start), - ("created_at_end", self.created_at_end), - ): - validate_filter_timestamp(value=value, field_name=field_name) - if ( - self.created_at_start is not None - and self.created_at_end is not None - and self.created_at_start > self.created_at_end - ): - raise InvalidFilterError("created_at_start must not be after created_at_end") + validate_history_run_filter(self) @dataclass(frozen=True) @@ -147,5 +101,4 @@ class RunPage: has_more: bool def __post_init__(self) -> None: - if not self.records and self.next_cursor is not None: - raise ProjectionConsistencyError("an empty run page cannot have a next cursor") + validate_history_page(page=self, is_event_page=False) diff --git a/src/sqlbuild/sqlite_history.py b/src/sqlbuild/sqlite_history.py index 121ddd7acf..e368a5e31c 100644 --- a/src/sqlbuild/sqlite_history.py +++ b/src/sqlbuild/sqlite_history.py @@ -1,11 +1,9 @@ """Typed public SQLite execution history API.""" -from sqlbuild.runtime.execution_history.classes.sqlite_execution_history import ( - SQLiteExecutionHistory as _SQLiteExecutionHistory, -) +from sqlbuild._native import SQLiteExecutionHistory as _NativeSQLiteExecutionHistory __all__ = ("SQLiteExecutionHistory",) -class SQLiteExecutionHistory(_SQLiteExecutionHistory): +class SQLiteExecutionHistory(_NativeSQLiteExecutionHistory): """Supported project-local SQLite execution history backend.""" diff --git a/tests/integration/src/sqlbuild/runtime/execution_history/test_sqlite_execution_history.py b/tests/integration/src/sqlbuild/runtime/execution_history/test_sqlite_execution_history.py index 8f4110dbe4..765a65665a 100644 --- a/tests/integration/src/sqlbuild/runtime/execution_history/test_sqlite_execution_history.py +++ b/tests/integration/src/sqlbuild/runtime/execution_history/test_sqlite_execution_history.py @@ -11,7 +11,6 @@ import pytest -import sqlbuild.runtime.execution_history.classes.sqlite_execution_history as sqlite_history_module from sqlbuild.execution_history import ( EventFilter, EventPage, @@ -48,16 +47,21 @@ def __init__( *, fail_commit_once: bool = False, fail_rollback_once: bool = False, + fail_projection_once: bool = False, ) -> None: self._connection: sqlite3.Connection = connection self._fail_commit_once: bool = fail_commit_once self._fail_rollback_once: bool = fail_rollback_once + self._fail_projection_once: bool = fail_projection_once @property def in_transaction(self) -> bool: return self._connection.in_transaction def execute(self, sql: str, parameters: Any = ()) -> sqlite3.Cursor: + if sql == "DELETE FROM run_projection" and self._fail_projection_once: + self._fail_projection_once = False + raise sqlite3.OperationalError("controlled body failure") if sql == "COMMIT" and self._fail_commit_once: self._fail_commit_once = False raise sqlite3.OperationalError("controlled commit failure") @@ -317,12 +321,14 @@ def test_given_body_and_rollback_failures_when_transaction_exits_then_body_error storage: SQLiteExecutionHistory = SQLiteExecutionHistory(path=tmp_path / "history.sqlite3") storage._connection = cast( sqlite3.Connection, - _TransactionFailureConnection(storage._connection, fail_rollback_once=True), + _TransactionFailureConnection( + storage._connection, fail_rollback_once=True, fail_projection_once=True + ), ) - with pytest.raises(ValueError, match="controlled body failure"): - with storage._transaction(): - raise ValueError("controlled body failure") + with pytest.raises(ExecutionHistoryStorageError) as raised: + storage.append_and_project((lifecycle_event("failed-body"),)) + assert str(raised.value.__cause__) == "controlled body failure" _ = storage.append_and_project((lifecycle_event("after-rollback-failure"),)) events: EventPage = storage.get_events(event_filter=EventFilter()) @@ -607,7 +613,7 @@ def tracked_connect( opened.append(connection) return connection - monkeypatch.setattr(sqlite_history_module.sqlite3, "connect", tracked_connect) + monkeypatch.setattr(sqlite3, "connect", tracked_connect) monkeypatch.setattr( SQLiteExecutionHistory, "_migrate", From 07d163188d1aeea9bdcc50f4f7836eb6f6744ae1 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Thu, 15 Oct 2026 13:41:29 +0100 Subject: [PATCH 12/12] fix: apply SQLite history file modes only on Unix --- .../_helpers/execution_history/sqlite.rs | 51 +++++++++++++++---- 1 file changed, 41 insertions(+), 10 deletions(-) diff --git a/crates/sqlbuild-python/src/bindings/_helpers/execution_history/sqlite.rs b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/sqlite.rs index 9148d2a76d..06767860f5 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/execution_history/sqlite.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/sqlite.rs @@ -1,8 +1,7 @@ //! `SQLiteExecutionHistory`: project-local SQLite event-log and run-projection storage. -use std::fs::{DirBuilder, Permissions}; +use std::fs::DirBuilder; use std::io::ErrorKind; -use std::os::unix::fs::{DirBuilderExt, PermissionsExt}; use std::path::Path; use std::sync::OnceLock; use std::sync::atomic::{AtomicBool, AtomicI64, Ordering}; @@ -13,12 +12,12 @@ use pyo3::types::{PyDict, PyTuple}; use sqlbuild_state::execution_history::constants::{ DEFAULT_PAGE_LIMIT, MEMORY_PATH, RUN_CURSOR_TIMESTAMP, SCHEMA_VERSION, SQLITE_ALL_EVENTS, SQLITE_ALL_RUNS, SQLITE_BEGIN, SQLITE_BUSY_TIMEOUT_MS, SQLITE_COMMIT, SQLITE_DELETE_RUNS, - SQLITE_DIRECTORY_MODE, SQLITE_EVENT_BY_ID, SQLITE_EVENT_BY_STORAGE_ID, SQLITE_EVENT_EXISTS, - SQLITE_FILE_MODE, SQLITE_FOREIGN_KEYS, SQLITE_HEALTHY_CHECK_RESULT, SQLITE_INSERT_NAMESPACE, - SQLITE_INSERT_SCHEMA_VERSION, SQLITE_JOURNAL_MODE, SQLITE_METADATA_EXISTS, - SQLITE_METADATA_TABLE, SQLITE_METADATA_VALUE, SQLITE_NAMESPACE_KEY, SQLITE_QUICK_CHECK, - SQLITE_ROLLBACK, SQLITE_RUN_BY_ID, SQLITE_RUN_EXISTS, SQLITE_RUN_INSERT, SQLITE_SCHEMA, - SQLITE_SCHEMA_VERSION_KEY, UNKNOWN_EVENT_CURSOR, UNKNOWN_RUN_CURSOR, + SQLITE_EVENT_BY_ID, SQLITE_EVENT_BY_STORAGE_ID, SQLITE_EVENT_EXISTS, SQLITE_FOREIGN_KEYS, + SQLITE_HEALTHY_CHECK_RESULT, SQLITE_INSERT_NAMESPACE, SQLITE_INSERT_SCHEMA_VERSION, + SQLITE_JOURNAL_MODE, SQLITE_METADATA_EXISTS, SQLITE_METADATA_TABLE, SQLITE_METADATA_VALUE, + SQLITE_NAMESPACE_KEY, SQLITE_QUICK_CHECK, SQLITE_ROLLBACK, SQLITE_RUN_BY_ID, SQLITE_RUN_EXISTS, + SQLITE_RUN_INSERT, SQLITE_SCHEMA, SQLITE_SCHEMA_VERSION_KEY, UNKNOWN_EVENT_CURSOR, + UNKNOWN_RUN_CURSOR, }; use sqlbuild_state::execution_history::errors::HistoryErrorKind; use sqlbuild_state::execution_history::main::decoded_run_cursor::decoded_run_cursor; @@ -77,18 +76,50 @@ fn prepare_history_directory(path: &str) -> std::io::Result<()> { { std::fs::create_dir_all(grandparent)?; } - match DirBuilder::new().mode(SQLITE_DIRECTORY_MODE).create(parent) { + match private_directory_builder().create(parent) { Err(error) if error.kind() == ErrorKind::AlreadyExists && parent.is_dir() => Ok(()), result => result, } } +/// A directory builder with the owner-only mode Python's `os.mkdir(mode=0o700)` applies. +#[cfg(unix)] +fn private_directory_builder() -> DirBuilder { + use std::os::unix::fs::DirBuilderExt; + let mut builder: DirBuilder = DirBuilder::new(); + builder.mode(sqlbuild_state::execution_history::constants::SQLITE_DIRECTORY_MODE); + builder +} + +/// Windows ignores the directory mode, as `os.mkdir` does there. +#[cfg(not(unix))] +fn private_directory_builder() -> DirBuilder { + DirBuilder::new() +} + /// `os.chmod(path, 0o600)` for a file-backed history. fn restrict_history_file(path: &str) -> std::io::Result<()> { if path == MEMORY_PATH { return Ok(()); } - std::fs::set_permissions(path, Permissions::from_mode(SQLITE_FILE_MODE)) + set_owner_only(path) +} + +#[cfg(unix)] +fn set_owner_only(path: &str) -> std::io::Result<()> { + use std::os::unix::fs::PermissionsExt; + std::fs::set_permissions( + path, + std::fs::Permissions::from_mode( + sqlbuild_state::execution_history::constants::SQLITE_FILE_MODE, + ), + ) +} + +/// `os.chmod(path, 0o600)` on Windows only keeps the file writable, which it already is. +#[cfg(not(unix))] +fn set_owner_only(_path: &str) -> std::io::Result<()> { + Ok(()) } /// Whether `error` is a `sqlite3.Error`.