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/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-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/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-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-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/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/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 f2911bf659..0a7ad4ba94 100644 --- a/crates/sqlbuild-core/src/lib.rs +++ b/crates/sqlbuild-core/src/lib.rs @@ -9,6 +9,8 @@ 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; pub mod cost_allocation; @@ -19,6 +21,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/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..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,13 +1,15 @@ //! 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, 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; @@ -17,6 +19,9 @@ 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 +87,70 @@ 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 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)) } -/// 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/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 aa0308a836..c09f1aea65 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, 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<()> { @@ -27,6 +27,9 @@ 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)?; + execution_history::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/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.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/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/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/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/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..a722957f46 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/compute_logs/storage.rs @@ -0,0 +1,460 @@ +//! `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 sqlbuild_cli::compute_logs::models::CaptureDocument; + +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: &[File; 3]) -> std::io::Result<[i64; 3]> { + let mut counts: [i64; 3] = [0; 3]; + for (count, mut stream) in counts.iter_mut().zip(streams.iter()) { + *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 written: Option> = { + let mut state = self.locked()?; + state + .writers + .get_mut(invocation_id) + .map(|writers| writers[stream].write_all(data)) + }; + 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]> { + 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 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. + #[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 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", + )), + }, + ) + }; + let (directory, streams) = checked(py, checked_streams)?; + 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/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..55a258af72 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_command/init_project.rs @@ -0,0 +1,350 @@ +//! 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::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_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: 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.")?; + 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, &PyString::new(py, &metadata.target_path).into_any())?, + )?; + 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")?, + &[ + ("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: 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(()) +} + +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: &DbtProjectMetadata, +) -> PyResult> { + let requested: Bound<'py, PyAny> = request.getattr("profile_name")?; + if requested.is_truthy()? { + return Ok(requested); + } + Ok(PyString::new(request.py(), &metadata.profile_name).into_any()) +} + +fn selected_output<'py>( + paths: &InitPaths<'py>, + request: &Bound<'py, PyAny>, + profile_name: &Bound<'py, PyAny>, +) -> PyResult> { + let target_name: Option = request.getattr("target_name")?.extract()?; + load_selected_dbt_profile_output( + &paths.profiles_dir, + &profile_name.extract::()?, + target_name.as_deref(), + ) +} + +/// `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..911be9bf38 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/dbt_interop/event_stream.rs @@ -0,0 +1,438 @@ +//! 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", + ], +); +/// 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); +/// 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, +} + +/// 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(PyErr::from)?; + waited = true; + py.check_signals()?; + Ok((returncode, results)) + })(); + let process_id: u32 = cleanup(py, &shared, refresher, (child, waited)); + let (returncode, results) = result?; + Ok((returncode, results, Some(process_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>, &[ActiveNode]), +) -> 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), +) -> u32 { + 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.id(); + } + py.detach(|| stop_dbt_process(child, terminate)).id() +} + +/// 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..79b76ab4ca 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; @@ -11,8 +12,10 @@ 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; 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/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/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/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..06767860f5 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/execution_history/sqlite.rs @@ -0,0 +1,952 @@ +//! `SQLiteExecutionHistory`: project-local SQLite event-log and run-projection storage. + +use std::fs::DirBuilder; +use std::io::ErrorKind; +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_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; +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 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(()); + } + 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`. +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 a63b0b867b..3277ea6a3a 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; @@ -31,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; @@ -44,6 +46,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/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..ca56e429af --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_evaluation.rs @@ -0,0 +1,511 @@ +//! 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::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}; +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::{ + HostCancel, HostRequest, host_payload_project, 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, + pub(crate) cancel: HostCancel, +} + +/// 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, + cancel: &run.cancel, + }, + &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) +} + +/// 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 new file mode 100644 index 0000000000..93fe2b427e --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs @@ -0,0 +1,488 @@ +//! 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::{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, 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; +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; +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"; +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", +]; + +/// 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>, + 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, + pub(crate) cancel: &'a HostCancel, +} + +/// 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 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, min_invocations_per_host)), + ) + }; + 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, 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, request.cancel), + (publish, &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, + launches: Option, +} + +impl Launch { + 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(); + 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, 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, + (launch, vec![spec]), + (write_inputs, abandoned), + (None, cancel), + )?; + 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, + (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]); + } + 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(|| 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<'_>, + (launch, specs): (&Launch, Vec), + (write_inputs, abandoned): (impl FnOnce() -> PyResult<()>, &Path), + (cancel_marker, cancel): (Option<&Path>, &HostCancel), +) -> 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 = 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); + 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 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() + }); + 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)?; + 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. +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(); + 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..4b10fcadae 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/mod.rs @@ -1,6 +1,9 @@ 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; pub(crate) mod dsl_headers; @@ -25,6 +28,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..efb60e8500 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/registration.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/registration.rs @@ -63,5 +63,8 @@ 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::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/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/rules/rule_dispatch.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/rule_dispatch.rs new file mode 100644 index 0000000000..04271a8ec9 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/rule_dispatch.rs @@ -0,0 +1,695 @@ +//! 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::{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}; + +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, +}; +use crate::bindings::_helpers::rules::custom_hosts::HostCancel; + +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"; +/// 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", + "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 = 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> { + 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>, + cancel: HostCancel, +} + +/// 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, + cancel: HostCancel::default(), + }; + 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), + cancel: HostCancel::default(), + }; + 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()? + }, + cancel: HostCancel::default(), + }) +} + +/// 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, + cancel: HostCancel::default(), + }; + 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_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) + }) + })?; + 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")?), + ], + ); + 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( + &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, &run.cancel); + } + 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,))); + cancel_on_interrupt(py, &evaluated, &run.cancel); + let custom: PyResult> = match custom { + Some(handle) => joined(py, handle, &run.cancel), + 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; + 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) +} + +/// 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> { + 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-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/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..d8da25d7de --- /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; + +/// 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, min_invocations_per_host): (usize, usize)) -> usize { + MAX_HOSTS + .min(cores) + .min(invocations / min_invocations_per_host.max(1)) + .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..f35430e473 --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/tests/test_custom_host_pool.rs @@ -0,0 +1,315 @@ +use std::collections::{BTreeMap, BTreeSet}; +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; +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, + 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, test_case.min_invocations_per_host) + ), + 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..16ea800bd3 --- /dev/null +++ b/crates/sqlbuild-rules/src/custom_host_pool/tests/test_types.rs @@ -0,0 +1,40 @@ +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) min_invocations_per_host: 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/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 15a559828f..bf472a55e7 100644 --- a/fensu.toml +++ b/fensu.toml @@ -43,6 +43,16 @@ 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" +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" @@ -68,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 1a723fac71..ec3df3e61e 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 = 9546 [[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" @@ -250,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 = "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 = "Documented SQLite and Postgres execution-history backends." +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" @@ -260,9 +255,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" @@ -270,9 +300,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" @@ -294,25 +354,10 @@ 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" -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" @@ -324,21 +369,11 @@ 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" 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" @@ -434,16 +469,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 +514,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 +529,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/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 c41205c948..6afbf93f48 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 @@ -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 @@ -258,9 +259,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 +274,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 @@ -307,14 +297,17 @@ 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, + DbtLsResult, DbtNodeExecutionResult, NormalizedDbtProfileConnection, ResolvedDbtProfileOutput, ) +from sqlbuild.integrations.dbt.types import DbtInvoker from sqlbuild.lint.models import ( CompileFacts, ExpandedLintFile, @@ -336,9 +329,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, @@ -347,10 +342,38 @@ 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.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 LifecycleEvent +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, @@ -3075,6 +3098,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, ...]: ... @@ -3411,6 +3470,179 @@ class CliObservabilityScope: def __exit__(self, *exc_info: object) -> bool: ... 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.""" + + 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: ... + +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: @@ -5172,8 +5404,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]: ... @@ -6450,38 +6681,73 @@ def custom_rule_identity( dialect: str, project_dir: str, ) -> str: ... +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 custom_host_payload(project: CompiledProjectView) -> CompiledProject: ... -class CustomRuleCache: - """One custom-rule run's subjects, reused and new results, and its cache file.""" - +class CustomHostObserver: 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]], - ] - ], + cores: int | None = None, + min_invocations_per_host: int | None = None, + publish_inputs: Callable[[Path], object] | None = None, ) -> None: ... - def findings(self) -> list[Finding]: ... - @property - def hits(self) -> int: ... - @property - def misses(self) -> int: ... + def launches(self) -> tuple[str, ...]: ... + +custom_host_observer: CustomHostObserver | None +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/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/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/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/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/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/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/_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/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/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/rule_engine/_helpers/engine/custom_dispatch.py b/src/sqlbuild/rule_engine/_helpers/engine/custom_dispatch.py deleted file mode 100644 index ae211e78c6..0000000000 --- a/src/sqlbuild/rule_engine/_helpers/engine/custom_dispatch.py +++ /dev/null @@ -1,178 +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.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 -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: - _ = 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 5acdaa97b9..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 - -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.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 -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=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/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/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/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/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/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/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/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 90% 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 index 014f8d7db2..f20ef69ab9 100644 --- a/src/sqlbuild/cli/commands/main/entrypoint/_dispatch_with_output_capture.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/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) 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/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/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/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_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 158d3ee9d7..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,6 +3,7 @@ from __future__ import annotations import json +import time from collections.abc import Callable from pathlib import Path from threading import Event @@ -13,15 +14,15 @@ 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 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, @@ -36,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, ) @@ -1032,7 +1037,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 +1045,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 +1102,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) - + 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) - assert len(host_runs) == 1 + assert len(hosts.launches()) == 1 assert warm == cold assert { f"models/orders.sql:{test_case.expected_code}", @@ -1158,31 +1156,27 @@ 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) - - 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) + 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) - monkeypatch.setattr(native_module, "run_custom_host_json", recording_host) - single: dict[str, Any] = compile_payload(hosts=1) - single_runs: int = len(host_runs) + 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(hosts=test_case.hosts) - split_runs: int = len(host_runs) - single_runs - warm: dict[str, Any] = compile_payload(hosts=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_runs, split_runs, len(host_runs)) == (1, test_case.hosts, 1 + test_case.hosts) + 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 // 3 assert warm["diagnostics"] == single["diagnostics"] @@ -1192,6 +1186,91 @@ def compile_payload(*, hosts: 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/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/integration/src/sqlbuild/rule_engine/host_payload/helpers.py b/tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py index f31880b78a..8419916287 100644 --- a/tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py +++ b/tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py @@ -1,20 +1,21 @@ 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 custom_rules, fact_replay +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._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 @@ -112,13 +113,14 @@ def fact_digests(*, project: CompiledProjectView, root: Path) -> dict[FactKey, t def record_trimmed_fields( monkeypatch: pytest.MonkeyPatch, ) -> dict[type[object], set[str]]: - """Record replaced fields and omitted constructor fields in the host-payload trimmers.""" + """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 dataclasses.replace(obj, **changes) + return replace(obj, **changes) def recording_model(**values: Any) -> CompiledModel: trimmed[CompiledModel].update( @@ -126,7 +128,6 @@ def recording_model(**values: Any) -> CompiledModel: ) 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) + 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 index b6ad426c3f..5088abeffd 100644 --- 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 @@ -6,11 +6,10 @@ 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._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 ( @@ -114,7 +113,8 @@ def test_given_compiled_project_when_trimming_host_payload_then_every_fact_is_un ) -> 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)) + 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) 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", 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/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/compiler/pipeline/main/test_compile_phase_timings.py b/tests/unit/src/sqlbuild/compiler/pipeline/main/test_compile_phase_timings.py index f580a17fd0..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,7 +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.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 @@ -80,7 +79,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/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_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" 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}'" 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..26da58fa3c 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/helpers.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/helpers.py @@ -9,13 +9,15 @@ import pytest -from sqlbuild import _native +import sqlbuild._native as native from sqlbuild._native import ( NativeRulesRequest, build_rule_catalogue, 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, @@ -27,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 -from sqlbuild.rule_engine.main._evaluate import evaluate from sqlbuild.rule_engine.models import ( CustomRulesOutcome, Finding, @@ -54,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: @@ -98,10 +97,10 @@ 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( + native.evaluate_native_rules( project=project, config=RulesConfig(cache=RulesCacheConfig(enabled=False)), project_dir=project_dir, @@ -383,7 +382,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/_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 9c581503f4..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,8 +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 @@ -14,12 +16,12 @@ import sqlbuild._native as native_module from sqlbuild._native import ( + CustomHostObserver, 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,32 +114,81 @@ def evaluate_order_rules( ) -def use_hosts(*, monkeypatch: pytest.MonkeyPatch, hosts: int) -> list[str]: +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.""" - launches: list[str] = [] - launch: Callable[[str], str] = native_module.run_custom_host_json + 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 recording_launch(spec_json: str) -> str: - launches.append(spec_json) - return launch(spec_json) +def host_inputs_left(root: Path) -> tuple[Path, ...]: + """Return every payload, marker or folder a run left in the host-input folder.""" - 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 + return tuple((root / "target" / "rules-cache" / "host-inputs").iterdir()) -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`.""" +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 - 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 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]]: @@ -442,47 +493,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..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 @@ -5,9 +5,9 @@ import pytest -from sqlbuild._native import rules_bulk_cache_path +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, custom_host_pool +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 @@ -17,14 +17,17 @@ HostCancellationTestCase, HostFailureTestCase, HostInputsTestCase, + HostInterruptTestCase, HostPartitionTestCase, - HostPlanTestCase, ) from tests.unit.src.sqlbuild.rule_engine._helpers.host.helpers import ( - delay_inputs, + child_processes, + delayed_publish, evaluate_order_rules, - fail_inputs, - fail_inputs_without_directory, + failed_publish, + failed_publish_without_directory, + host_inputs_left, + interrupting_publish, orders_project, record_signals, touch_files, @@ -82,21 +85,21 @@ def test_given_partitioned_hosts_when_evaluating_then_findings_and_cache_match_o 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_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) - monkeypatch.undo() - split_launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) + 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_launches) + 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 @@ -104,7 +107,7 @@ def test_given_partitioned_hosts_when_evaluating_then_findings_and_cache_match_o assert split.findings == single.findings assert split_warm.findings == single_warm.findings == single.findings - assert len(single_launches) == 1 + int(not cache_enabled) + 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) == ( @@ -113,7 +116,7 @@ def test_given_partitioned_hosts_when_evaluating_then_findings_and_cache_match_o ) assert split_cache == single_cache assert bool(single_cache) is cache_enabled - assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) + assert not host_inputs_left(tmp_path) @pytest.mark.parametrize( @@ -133,12 +136,12 @@ def test_given_results_cached_by_split_hosts_when_rerunning_then_only_edited_sub 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) + hosts: CustomHostObserver = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) cold: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) - cold_launches: int = len(launches) + cold_launches: int = len(hosts.launches()) warm: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) - warm_launches: int = len(launches) - cold_launches + 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 @@ -179,52 +182,13 @@ def test_given_failing_rule_on_split_hosts_when_evaluating_then_single_host_erro _ = 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 + assert not host_inputs_left(tmp_path) @pytest.mark.parametrize( @@ -234,7 +198,6 @@ def test_given_planned_invocations_when_partitioning_then_split_is_deterministic 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, @@ -247,8 +210,8 @@ def test_given_failing_host_when_others_are_busy_then_they_stop_before_later_inv ) -> 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) + 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: @@ -256,38 +219,45 @@ def test_given_failing_host_when_others_are_busy_then_they_stop_before_later_inv 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()) + assert len(hosts.launches()) == test_case.hosts + 1 + assert signals == [] + assert not host_inputs_left(tmp_path) @pytest.mark.parametrize( "test_case", [ - HostCancellationTestCase( - description="busy hosts stop at their next invocation", + HostInterruptTestCase( + description="busy split hosts stop after an interrupt", 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, + delay_seconds=1.0, + interrupt_after_seconds=1.5, + expected_max_seconds=10.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 +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, failing=("orders_000",), delay_seconds=test_case.delay_seconds) + write_order_rules(root=tmp_path, 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) + 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(RulesError) as raised: + with pytest.raises(KeyboardInterrupt): _ = evaluate_order_rules(project=project, root=tmp_path) - assert signals == [] - assert test_case.expected_error_fragment in str(raised.value) + 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( @@ -318,16 +288,19 @@ def test_given_hosts_started_before_payload_when_payload_arrives_then_findings_a 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) + 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(launches) == test_case.hosts + assert len(hosts.launches()) == test_case.hosts assert delayed.findings == oracle.findings - assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) + assert not host_inputs_left(tmp_path) @pytest.mark.parametrize( @@ -355,16 +328,19 @@ def test_given_payload_write_failure_when_hosts_wait_then_original_error_is_rais ) -> 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) + 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(launches) == test_case.hosts - assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) + assert len(hosts.launches()) == test_case.hosts + assert not host_inputs_left(tmp_path) @pytest.mark.parametrize( @@ -392,17 +368,17 @@ def test_given_input_folder_removed_when_payload_write_fails_then_original_error ) -> 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 + 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(launches) == test_case.hosts - assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) + assert len(hosts.launches()) == test_case.hosts + assert not host_inputs_left(tmp_path) @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 4a74d73192..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,25 +2,12 @@ import pytest +import sqlbuild._native as native 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 -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.""" @@ -39,5 +26,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") 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 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", ( 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, )