diff --git a/AGENTS.md b/AGENTS.md index 96562afb4d..30e75b1806 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -66,6 +66,19 @@ affects analysis, validation, or caching must report all three. - Performance fixes must not reduce coverage. Show that diagnostics are unchanged or a strict, validated superset on the dense benchmark and a real-project run. +- When load makes wall times unreliable, compare `perf stat -e instructions:u` (repeat runs agree + to about 0.05%). Instruction counts cannot show overlap, so judge parallelism and background-work + changes by paired wall time from quiet runs. +- Before timing a candidate, prove its extension binary contains the change (for example, a new + symbol or string in the built `.so`); a stale build silently measures the base. +- Read the timeline for main-thread/pool alternation: serial native loops on the main thread while + the pool idles cost more wall time than CPU-heavy phases that already use every core. +- Work that runs concurrently with another phase saves wall time only while it is the longer of the + two; identify the critical path before optimising one side. Report the overlapped phase's own + timing too, because background work can slow the phase it overlaps by as much as it saves. +- Speculative early work must verify or discard: compare a digest of the real request at the normal + call site, hold cache writes until adoption, never run user code early, and treat start, panic + or evaluation failures as misses that fall back to the normal path. - Do not raise performance budgets or tolerances to make a regression pass. - Temporary user-approved exception (2026-10-10), removable after post-conversion memory optimisation: only the 10,000-model case in `test_fresh_process_compile_cache_performance.py` diff --git a/crates/sqlbuild-python/src/bindings/_helpers/model_config/header_metadata.rs b/crates/sqlbuild-python/src/bindings/_helpers/model_config/header_metadata.rs index b99f80b25d..250844aba9 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/model_config/header_metadata.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/model_config/header_metadata.rs @@ -176,14 +176,47 @@ impl HeaderColumn { } } +impl HeaderColumn { + /// The natively parsed column, or `None` when Python built its values or record. + pub(crate) fn native_spec(&self) -> Option<&NativeColumnSpec> { + match &self.fields { + ColumnFields::Native(spec) => Some(spec), + ColumnFields::Python(_) | ColumnFields::Projected => None, + } + } +} + +/// The `SchemaColumn` a native column projects to when its location names `location`'s path. +pub(crate) fn native_column_record<'py>( + py: Python<'py>, + column: &NativeColumn, + location: Option<(Bound<'py, PyAny>, LineColumnSpan)>, +) -> PyResult> { + let location: Bound<'py, PyAny> = match location { + Some((path, span)) => location_record(py, path, span)?, + None => py.None().into_bound(py), + }; + contract_object( + &record_class(py, RecordId::SchemaColumn)?, + column_values(py, column, location)?, + ) +} + /// A native column's record fields, in declaration order. fn native_values<'py>( py: Python<'py>, spec: &NativeColumnSpec, ) -> PyResult<[Bound<'py, PyAny>; 8]> { - let location: Bound<'py, PyAny> = native_location(py, spec)?; - let audits: Bound<'py, PyAny> = native_audits(py, spec, &location)?; - let column: &NativeColumn = &spec.column; + column_values(py, &spec.column, native_location(py, spec)?) +} + +/// A native column's record fields at `location`, in declaration order. +fn column_values<'py>( + py: Python<'py>, + column: &NativeColumn, + location: Bound<'py, PyAny>, +) -> PyResult<[Bound<'py, PyAny>; 8]> { + let audits: Bound<'py, PyAny> = native_audits(py, column, &location)?; Ok([ column.name.as_str().into_bound_py_any(py)?, column.column_type.as_deref().into_bound_py_any(py)?, @@ -209,7 +242,7 @@ fn native_value<'py>( "nullable" => column.nullable.into_bound_py_any(py), "description" => column.description.as_deref().into_bound_py_any(py), "meta" => Ok(PyDict::new(py).into_any()), - "audits" => native_audits(py, spec, &native_location(py, spec)?), + "audits" => native_audits(py, column, &native_location(py, spec)?), "location" => native_location(py, spec), _ => column.migrate_from.as_deref().into_bound_py_any(py), } @@ -218,31 +251,35 @@ fn native_value<'py>( /// The column's `SourceLocation`, or `None`. fn native_location<'py>(py: Python<'py>, spec: &NativeColumnSpec) -> PyResult> { match &spec.location { - Some((model_file, span)) => new_record( - py, - RecordId::SourceLocation, - ( - model_file.get().path_objects(py)?.1, - span.line, - span.column, - span.end_line, - span.end_column, - ), - None, - ), + Some((model_file, span)) => { + location_record(py, model_file.get().path_objects(py)?.1, *span) + } None => Ok(py.None().into_bound(py)), } } +/// `SourceLocation(path, line, column, end_line, end_column)`. +fn location_record<'py>( + py: Python<'py>, + path: Bound<'py, PyAny>, + span: LineColumnSpan, +) -> PyResult> { + new_record( + py, + RecordId::SourceLocation, + (path, span.line, span.column, span.end_line, span.end_column), + None, + ) +} + /// The column's audits as a tuple of `SchemaAuditInstance`s at `location`. fn native_audits<'py>( py: Python<'py>, - spec: &NativeColumnSpec, + column: &NativeColumn, location: &Bound<'py, PyAny>, ) -> PyResult> { let classes: ContractClasses<'py> = contract_classes(py)?.extract()?; - let audits: Vec> = spec - .column + let audits: Vec> = column .audits .iter() .map(|audit| native_audit_instance(py, &classes, audit, location)) diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_evaluation.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_evaluation.rs index d6b9e3722e..f5169cf5c5 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_evaluation.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_evaluation.rs @@ -5,7 +5,7 @@ use std::path::{Path, PathBuf}; use std::time::Instant; use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyModule, PyModuleMethods, PyResult, Python}; -use pyo3::types::{PyDict, PyDictMethods, PyTuple}; +use pyo3::types::{PyBytes, PyTuple}; use pyo3::{IntoPyObjectExt, pyfunction, wrap_pyfunction}; use sqlbuild_cache::digest::main::sha256_hex::sha256_hex; use sqlbuild_core::json::main::dumps::dumps; @@ -18,17 +18,17 @@ 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::imported; use crate::bindings::_helpers::compile_driver::python_calls::keywords; -use crate::bindings::_helpers::compiled_project::records::compiled_project_record; -use crate::bindings::_helpers::records::construct::{new_record, new_record_with}; +use crate::bindings::_helpers::records::construct::new_record_with; use crate::bindings::_helpers::records::fields::field as native_field; -use crate::bindings::_helpers::records::fields::replaced_keywords; use crate::bindings::_helpers::records::ids::RecordId; use crate::bindings::_helpers::rules::constants::{ CUSTOM_HOST_RUNTIME_VERSION, CUSTOM_RULES_CACHE_FILE, }; -use crate::bindings::_helpers::rules::custom_hosts::{ - HostCancel, HostRequest, host_payload_project, run_custom_hosts, +use crate::bindings::_helpers::rules::custom_host_columns::custom_host_declared_columns; +use crate::bindings::_helpers::rules::custom_host_payload::{ + CUSTOM_HOST_PAYLOAD_MAGIC, custom_host_payload, }; +use crate::bindings::_helpers::rules::custom_hosts::{HostCancel, HostRequest, run_custom_hosts}; use crate::bindings::_helpers::rules::custom_rule_cache::CustomRuleCache; use crate::bindings::_helpers::rules::custom_rule_cache::custom_rule_identity; use crate::bindings::_helpers::rules::custom_rule_evidence::custom_rule_implementation_fingerprint; @@ -39,29 +39,6 @@ use crate::bindings::_helpers::rules::rule_run::rules_bulk_cache_path; 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>, @@ -140,11 +117,10 @@ pub(crate) fn evaluate_custom_rules_cached<'py>( .iter() .map(|(code, rule)| Ok((code.clone(), host_subject(rule)?))) .collect::>()?; - let facts_project: Bound<'py, PyAny> = host_project(run.project)?; let host_run: HostRun = run_custom_hosts( py, &HostRequest { - project: &facts_project, + project: run.project, config: run.config, project_dir: &project_dir, dialect: run.dialect, @@ -301,25 +277,6 @@ fn orjson_key(key: &FactKey) -> String { 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 record: Bound<'py, PyAny> = compiled_project_record(project, false)?; - 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", - new_record(py, RecordId::ScopeIndex, (), None)?, - )?; - replaced_keywords(&record, &changes) -} - fn subject_value(rule: &Bound<'_, PyAny>) -> PyResult> { let subject: Bound<'_, PyAny> = native_field(rule, "subject")?; if subject.is_none() { @@ -389,12 +346,6 @@ fn cpu_seconds(py: Python<'_>) -> PyResult { 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)?) -} - /// The digest the custom-rule cache keys one recorded fact on, or `None` when unreproducible. #[pyfunction] #[pyo3(signature = (*, project, project_dir, key))] @@ -408,5 +359,10 @@ fn custom_rule_fact_digest( pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(custom_host_payload, module)?)?; + module.add_function(wrap_pyfunction!(custom_host_declared_columns, module)?)?; + module.add( + "CUSTOM_HOST_PAYLOAD_MAGIC", + PyBytes::new(module.py(), CUSTOM_HOST_PAYLOAD_MAGIC), + )?; module.add_function(wrap_pyfunction!(custom_rule_fact_digest, module)?) } diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_columns.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_columns.rs new file mode 100644 index 0000000000..559ebfa4f5 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_columns.rs @@ -0,0 +1,254 @@ +//! Natively parsed header columns as custom-rule host payloads carry them. + +use std::path::PathBuf; + +use pyo3::prelude::{Bound, PyAny, PyResult, Python}; +use pyo3::types::PyTuple; +use pyo3::{IntoPyObject, pyfunction}; +use serde::{Deserialize, Serialize}; +use sqlbuild_discovery::model_files::models::{NativeAudit, NativeColumn}; +use sqlbuild_discovery::models::LineColumnSpan; +use sqlbuild_model_config::header_metadata::models::{ParsedThresholds, ThresholdBound}; +use sqlbuild_sqltext::compiler::models::AuthoredValue; + +use crate::bindings::_helpers::model_config::header_metadata::native_column_record; +use crate::bindings::_helpers::rules::custom_hosts::rules_error; + +/// One native header column with the text of the path its location names. +#[derive(Debug, Serialize, Deserialize)] +pub(crate) struct PayloadColumn { + name: String, + column_type: Option, + nullable: Option, + description: Option, + audits: Vec, + migrate_from: Option, + location: Option<(String, [usize; 4])>, +} + +#[derive(Debug, Serialize, Deserialize)] +struct PayloadAudit { + definition_name: String, + arguments: Vec<(String, PayloadValue)>, + name: Option, + description: Option, + severity: Option, + run_scope: Option, + always_run: Option, + thresholds: Option<(Option, Option)>, + minimum_samples: Option, + evidence_limit: Option, +} + +/// A threshold bound with its limits as `f64` bits, so every value survives exactly. +#[derive(Debug, Serialize, Deserialize)] +enum PayloadBound { + Below(u64), + Above(u64), + Outside(u64, u64), +} + +#[derive(Debug, Serialize, Deserialize)] +enum PayloadValue { + Null, + Boolean(bool), + BareWord(String), + String(String), + List(Vec), + Map(Vec<(String, PayloadValue)>), + Set(Vec), + Tuple(Vec), + TypedConstant(Vec<(String, PayloadValue)>), + InlineSqlHook(String), + NamedSqlHook(String, Vec<(String, PayloadValue)>), + PythonHook(String, Vec<(String, PayloadValue)>), +} + +impl PayloadColumn { + /// `column` with its location naming `path`. + pub(crate) fn new(column: &NativeColumn, location: Option<(String, LineColumnSpan)>) -> Self { + Self { + name: column.name.clone(), + column_type: column.column_type.clone(), + nullable: column.nullable, + description: column.description.clone(), + audits: column.audits.iter().map(PayloadAudit::new).collect(), + migrate_from: column.migrate_from.clone(), + location: location.map(|(path, span)| { + ( + path, + [span.line, span.column, span.end_line, span.end_column], + ) + }), + } + } + + fn column(self) -> (NativeColumn, Option<(String, LineColumnSpan)>) { + let location: Option<(String, LineColumnSpan)> = + self.location + .map(|(path, [line, column, end_line, end_column])| { + ( + path, + LineColumnSpan { + line, + column, + end_line, + end_column, + }, + ) + }); + let column: NativeColumn = NativeColumn { + name: self.name, + column_type: self.column_type, + nullable: self.nullable, + description: self.description, + audits: self.audits.into_iter().map(PayloadAudit::audit).collect(), + migrate_from: self.migrate_from, + }; + (column, location) + } +} + +impl PayloadAudit { + fn new(audit: &NativeAudit) -> Self { + Self { + definition_name: audit.definition_name.clone(), + arguments: pairs(&audit.arguments), + name: audit.name.clone(), + description: audit.description.clone(), + severity: audit.severity.clone(), + run_scope: audit.run_scope.clone(), + always_run: audit.always_run, + thresholds: audit.thresholds.map(|thresholds| { + ( + thresholds.warn.map(PayloadBound::new), + thresholds.error.map(PayloadBound::new), + ) + }), + minimum_samples: audit.minimum_samples.as_ref().map(PayloadValue::new), + evidence_limit: audit.evidence_limit.as_ref().map(PayloadValue::new), + } + } + + fn audit(self) -> NativeAudit { + NativeAudit { + definition_name: self.definition_name, + arguments: authored_pairs(self.arguments), + name: self.name, + description: self.description, + severity: self.severity, + run_scope: self.run_scope, + always_run: self.always_run, + thresholds: self.thresholds.map(|(warn, error)| ParsedThresholds { + warn: warn.map(PayloadBound::bound), + error: error.map(PayloadBound::bound), + }), + minimum_samples: self.minimum_samples.map(PayloadValue::authored), + evidence_limit: self.evidence_limit.map(PayloadValue::authored), + } + } +} + +impl PayloadBound { + fn new(bound: ThresholdBound) -> Self { + match bound { + ThresholdBound::Below(limit) => Self::Below(limit.to_bits()), + ThresholdBound::Above(limit) => Self::Above(limit.to_bits()), + ThresholdBound::Outside(lower, upper) => { + Self::Outside(lower.to_bits(), upper.to_bits()) + } + } + } + + fn bound(self) -> ThresholdBound { + match self { + Self::Below(limit) => ThresholdBound::Below(f64::from_bits(limit)), + Self::Above(limit) => ThresholdBound::Above(f64::from_bits(limit)), + Self::Outside(lower, upper) => { + ThresholdBound::Outside(f64::from_bits(lower), f64::from_bits(upper)) + } + } + } +} + +impl PayloadValue { + fn new(value: &AuthoredValue) -> Self { + match value { + AuthoredValue::Null => Self::Null, + AuthoredValue::Boolean(flag) => Self::Boolean(*flag), + AuthoredValue::BareWord(text) => Self::BareWord(text.clone()), + AuthoredValue::String(text) => Self::String(text.clone()), + AuthoredValue::List(items) => Self::List(items.iter().map(Self::new).collect()), + AuthoredValue::Map(entries) => Self::Map(pairs(entries)), + AuthoredValue::Set(items) => Self::Set(items.iter().map(Self::new).collect()), + AuthoredValue::Tuple(items) => Self::Tuple(items.iter().map(Self::new).collect()), + AuthoredValue::TypedConstant(entries) => Self::TypedConstant(pairs(entries)), + AuthoredValue::InlineSqlHook(text) => Self::InlineSqlHook(text.clone()), + AuthoredValue::NamedSqlHook(name, entries) => { + Self::NamedSqlHook(name.clone(), pairs(entries)) + } + AuthoredValue::PythonHook(name, entries) => { + Self::PythonHook(name.clone(), pairs(entries)) + } + } + } + + fn authored(self) -> AuthoredValue { + match self { + Self::Null => AuthoredValue::Null, + Self::Boolean(flag) => AuthoredValue::Boolean(flag), + Self::BareWord(text) => AuthoredValue::BareWord(text), + Self::String(text) => AuthoredValue::String(text), + Self::List(items) => { + AuthoredValue::List(items.into_iter().map(Self::authored).collect()) + } + Self::Map(entries) => AuthoredValue::Map(authored_pairs(entries)), + Self::Set(items) => AuthoredValue::Set(items.into_iter().map(Self::authored).collect()), + Self::Tuple(items) => { + AuthoredValue::Tuple(items.into_iter().map(Self::authored).collect()) + } + Self::TypedConstant(entries) => AuthoredValue::TypedConstant(authored_pairs(entries)), + Self::InlineSqlHook(text) => AuthoredValue::InlineSqlHook(text), + Self::NamedSqlHook(name, entries) => { + AuthoredValue::NamedSqlHook(name, authored_pairs(entries)) + } + Self::PythonHook(name, entries) => { + AuthoredValue::PythonHook(name, authored_pairs(entries)) + } + } + } +} + +fn pairs(entries: &[(String, AuthoredValue)]) -> Vec<(String, PayloadValue)> { + entries + .iter() + .map(|(key, value)| (key.clone(), PayloadValue::new(value))) + .collect() +} + +fn authored_pairs(entries: Vec<(String, PayloadValue)>) -> Vec<(String, AuthoredValue)> { + entries + .into_iter() + .map(|(key, value)| (key, value.authored())) + .collect() +} + +/// The `SchemaColumn`s of one model's payload columns, built as the compiling process built them. +#[pyfunction] +pub(crate) fn custom_host_declared_columns<'py>( + py: Python<'py>, + data: &[u8], +) -> PyResult> { + let columns: Vec = + serde_json::from_slice(data).map_err(|error| rules_error(py, error.to_string()))?; + let mut records: Vec> = Vec::with_capacity(columns.len()); + for column in columns { + let (column, location) = column.column(); + let location: Option<(Bound<'py, PyAny>, LineColumnSpan)> = match location { + Some((path, span)) => Some((PathBuf::from(path).into_pyobject(py)?.into_any(), span)), + None => None, + }; + records.push(native_column_record(py, &column, location)?); + } + PyTuple::new(py, records) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_payload.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_payload.rs new file mode 100644 index 0000000000..e18d3fda33 --- /dev/null +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_payload.rs @@ -0,0 +1,328 @@ +//! The custom-rule host payload: the facts rules read, written natively for hosts to decode lazily. + +use std::io::Write; +use std::sync::Arc; + +use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyResult, PyTypeMethods, Python}; +use pyo3::pyfunction; +use pyo3::types::{ + PyBytes, PyBytesMethods, PyDictMethods, PyString, PyStringMethods, PyTuple, PyTupleMethods, +}; +use serde_json::{Map, Value, json}; +use sqlbuild_discovery::models::LineColumnSpan; + +use crate::bindings::_helpers::compiled_project::compiled_model_view::NativeCompiledModel; +use crate::bindings::_helpers::compiled_project::records::{ + compiled_audit_record, compiled_function_record, compiled_seed_record, compiled_source_record, +}; +use crate::bindings::_helpers::records::fields::field as native_field; +use crate::bindings::_helpers::rules::custom_host_columns::PayloadColumn; +use crate::bindings::_helpers::rules::custom_hosts::pickled; +use crate::bindings::_helpers::rules::fact_replay::{ + ModelEntry, project_attribute, project_models, +}; + +/// The first bytes of every payload file; hosts reject any other. +pub(crate) const CUSTOM_HOST_PAYLOAD_MAGIC: &[u8; 8] = b"SQBHOST1"; +/// The config values model facts read. +const CONFIG_KEYS: [&str; 3] = ["materialized", "contract", "grain"]; +const PICKLE_TAG: &str = "p"; +const TEXT_TAG: &str = "s"; +const TEXTS_TAG: &str = "t"; +const INFERRED_TAG: &str = "r"; +const COLUMNS_TAG: &str = "c"; +const PATHLIB_MODULE: &str = "pathlib"; +const PATH_CLASSES: [&str; 2] = ["PosixPath", "WindowsPath"]; + +type RecordConversion = for<'py> fn(&Bound<'py, PyAny>) -> PyResult>; + +/// Project collections in header order, with the record each item becomes. +const SECTIONS: [(&str, Option); 7] = [ + ("sources", Some(compiled_source_record)), + ("seeds", Some(compiled_seed_record)), + ("functions", Some(compiled_function_record)), + ("audits", Some(compiled_audit_record)), + ("sql_tests", None), + ("public_enums", None), + ("public_constants", None), +]; + +/// What a model's declared columns are written as. +enum Declared { + Absent, + Native(Vec), + Python, +} + +/// Writes payload parts in order and remembers where each one lies. +struct PayloadWriter { + out: W, + offset: u64, +} + +/// Write `project`'s facts and `config` to `out`: magic, parts, JSON header, header offset. +pub(crate) fn write_payload( + out: W, + project: &Bound<'_, PyAny>, + config: &Bound<'_, PyAny>, +) -> PyResult { + let py: Python<'_> = project.py(); + let mut writer: PayloadWriter = PayloadWriter { out, offset: 0 }; + writer.write(CUSTOM_HOST_PAYLOAD_MAGIC)?; + let mut models: Vec = Vec::new(); + for entry in project_models(project)? { + models.push(writer.model(py, &entry)?); + } + let mut sections: Map = Map::new(); + for (name, conversion) in SECTIONS { + let value: Bound<'_, PyAny> = project_attribute(project, name)?; + let value: Bound<'_, PyAny> = match conversion { + Some(record) => PyTuple::new( + py, + value + .try_iter()? + .map(|item| record(&item?)) + .collect::>>()?, + )? + .into_any(), + None => value, + }; + sections.insert(name.to_owned(), writer.pickle(&value)?); + } + let header: Value = json!({ + "target": writer.scalar(&native_field(project, "effective_target_name")?)?, + "config": writer.pickle(config)?, + "sections": sections, + "models": models, + }); + let header_offset: u64 = writer.offset; + let text: Vec = serde_json::to_vec(&header) + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string()))?; + writer.write(&text)?; + writer.write(&header_offset.to_le_bytes())?; + Ok(writer.out) +} + +impl PayloadWriter { + fn write(&mut self, bytes: &[u8]) -> PyResult<()> { + self.out.write_all(bytes)?; + self.offset += bytes.len() as u64; + Ok(()) + } + + /// `bytes` as one part tagged `tag`. + fn part(&mut self, tag: &str, bytes: &[u8]) -> PyResult { + let start: u64 = self.offset; + self.write(bytes)?; + Ok(json!([tag, start, bytes.len()])) + } + + fn pickle(&mut self, value: &Bound<'_, PyAny>) -> PyResult { + let bytes: Bound<'_, PyBytes> = pickled(value)?; + self.part(PICKLE_TAG, bytes.as_bytes()) + } + + /// `None`, a `str`, or a tuple of `str` inline; any other value pickled. + fn scalar(&mut self, value: &Bound<'_, PyAny>) -> PyResult { + if value.is_none() { + return Ok(Value::Null); + } + if let Some(text) = exact_text(value)? { + return Ok(Value::String(text)); + } + if value.is_exact_instance_of::() { + let items: &Bound<'_, PyTuple> = value.downcast::()?; + let mut texts: Vec = Vec::with_capacity(items.len()); + for item in items.iter() { + let Some(text) = exact_text(&item)? else { + return self.pickle(value); + }; + texts.push(Value::String(text)); + } + return Ok(json!([TEXTS_TAG, texts])); + } + self.pickle(value) + } + + /// A `str` as a text part; any other value pickled. + fn text(&mut self, value: &Bound<'_, PyAny>) -> PyResult { + if let Some(text) = exact_text(value)? { + return self.part(TEXT_TAG, text.as_bytes()); + } + self.pickle(value) + } + + /// A concrete `pathlib.Path` as its text; any other value pickled. + fn path(&mut self, path: &Bound<'_, PyAny>) -> PyResult { + match plain_path_text(path)? { + Some(text) => Ok(Value::String(text)), + None => self.pickle(path), + } + } + + /// An empty tuple as `null`; any other value pickled. + fn declarations(&mut self, value: &Bound<'_, PyAny>) -> PyResult { + if value.is_exact_instance_of::() && value.downcast::()?.is_empty() { + return Ok(Value::Null); + } + self.pickle(value) + } + + fn model(&mut self, py: Python<'_>, entry: &ModelEntry<'_>) -> PyResult { + let mut model: Map = Map::new(); + model.insert("name".to_owned(), Value::String(entry.name.clone())); + model.insert("path".to_owned(), self.path(&entry.path)?); + model.insert("key".to_owned(), json!(entry.key)); + model.insert("deps".to_owned(), json!(entry.deps)); + let mut config: Map = Map::new(); + if let Some(values) = &entry.config_values { + for name in CONFIG_KEYS { + if let Some(value) = values.get_item(name)? { + config.insert(name.to_owned(), self.scalar(&value)?); + } + } + } + model.insert("config".to_owned(), Value::Object(config)); + let facts: Map = match entry.model.downcast::() { + Ok(view) => self.native_model(py, view.get())?, + Err(_) => self.record_model(&entry.model)?, + }; + model.extend(facts); + Ok(Value::Object(model)) + } + + /// A native view's facts, read from the compiled project without Python records. + fn native_model( + &mut self, + py: Python<'_>, + view: &NativeCompiledModel, + ) -> PyResult> { + let mut model: Map = Map::new(); + let (authored, query): (Arc, Arc) = view.read(py, |record, _| { + Ok(( + record.model_file.get().contents_text().clone(), + record.query_sql.clone(), + )) + })?; + model.insert( + "authored".to_owned(), + self.part(TEXT_TAG, authored.as_bytes())?, + ); + model.insert("query".to_owned(), self.part(TEXT_TAG, query.as_bytes())?); + let inferred: Value = match view.inferred_rows(py)? { + Some(rows) => self.part(INFERRED_TAG, &json_bytes(&rows)?)?, + None => Value::Null, + }; + model.insert("inferred".to_owned(), inferred); + let declared: Value = match native_declared(py, view)? { + Declared::Absent => Value::Null, + Declared::Native(columns) => self.part(COLUMNS_TAG, &json_bytes(&columns)?)?, + Declared::Python => match view.declared_columns(py)? { + Some(columns) => self.pickle(columns.bind(py).as_any())?, + None => Value::Null, + }, + }; + model.insert("declared".to_owned(), declared); + let enums: Bound<'_, PyAny> = view.enum_declarations(py)?.into_bound(py).into_any(); + model.insert("enums".to_owned(), self.declarations(&enums)?); + let constants: Bound<'_, PyAny> = view.constant_declarations(py)?.into_bound(py).into_any(); + model.insert("constants".to_owned(), self.declarations(&constants)?); + Ok(model) + } + + /// A model record's facts, read through its fields. + fn record_model(&mut self, record: &Bound<'_, PyAny>) -> PyResult> { + let mut model: Map = Map::new(); + let authored: Value = self.text(&native_field(record, "authored_sql")?)?; + model.insert("authored".to_owned(), authored); + let query: Value = self.text(&native_field(record, "query_sql")?)?; + model.insert("query".to_owned(), query); + let inferred: Bound<'_, PyAny> = native_field(record, "inferred_columns")?; + let inferred: Value = if inferred.is_none() { + Value::Null + } else { + self.pickle(&inferred)? + }; + model.insert("inferred".to_owned(), inferred); + let entry: Bound<'_, PyAny> = native_field(record, "schema_entry")?; + let declared: Value = if entry.is_none() { + Value::Null + } else { + self.pickle(&native_field(&entry, "columns")?)? + }; + model.insert("declared".to_owned(), declared); + let enums: Value = self.declarations(&native_field(record, "enum_declarations")?)?; + model.insert("enums".to_owned(), enums); + let constants: Value = + self.declarations(&native_field(record, "constant_declarations")?)?; + model.insert("constants".to_owned(), constants); + Ok(model) + } +} + +/// The view's declared columns when every one was parsed natively at a plain path. +fn native_declared(py: Python<'_>, view: &NativeCompiledModel) -> PyResult { + view.read(py, |record, py| { + let Some(schema) = &record.schema else { + return Ok(Declared::Absent); + }; + let mut columns: Vec = Vec::with_capacity(schema.columns.specs.len()); + for column in &schema.columns.specs { + let Some(spec) = column.native_spec() else { + return Ok(Declared::Python); + }; + let location: Option<(String, LineColumnSpan)> = match &spec.location { + Some((file, span)) => match plain_path_text(&file.get().path_objects(py)?.1)? { + Some(text) => Some((text, *span)), + None => return Ok(Declared::Python), + }, + None => None, + }; + columns.push(PayloadColumn::new(&spec.column, location)); + } + Ok(Declared::Native(columns)) + }) +} + +/// The text of a concrete `pathlib.Path`, which `Path(text)` rebuilds equal, or `None`. +fn plain_path_text(path: &Bound<'_, PyAny>) -> PyResult> { + let class = path.get_type(); + let module: String = class.module()?.to_str()?.to_owned(); + let qualname: String = class.qualname()?.to_str()?.to_owned(); + if module.split('.').next() != Some(PATHLIB_MODULE) + || !PATH_CLASSES.contains(&qualname.as_str()) + { + return Ok(None); + } + if let Ok(text) = path.str()?.to_str() { + return Ok(Some(text.to_owned())); + } + Ok(None) +} + +/// The text of an exact `str` that UTF-8 can carry, or `None`. +fn exact_text(value: &Bound<'_, PyAny>) -> PyResult> { + if !value.is_exact_instance_of::() { + return Ok(None); + } + if let Ok(text) = value.downcast::()?.to_str() { + return Ok(Some(text.to_owned())); + } + Ok(None) +} + +fn json_bytes(value: &T) -> PyResult> { + serde_json::to_vec(value) + .map_err(|error| pyo3::exceptions::PyValueError::new_err(error.to_string())) +} + +/// The payload a custom-rule host decodes for `project` and `config`. +#[pyfunction] +#[pyo3(signature = (project, config))] +pub(crate) fn custom_host_payload<'py>( + project: &Bound<'py, PyAny>, + config: &Bound<'py, PyAny>, +) -> PyResult> { + let bytes: Vec = write_payload(Vec::new(), project, config)?; + Ok(PyBytes::new(project.py(), &bytes)) +} diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs index 4026874001..7cabe5d7d5 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs @@ -1,6 +1,8 @@ //! Run planned custom-rule invocations across isolated host processes and merge their results. use std::collections::BTreeSet; +use std::fs::File; +use std::io::{BufWriter, Write}; use std::path::{Path, PathBuf}; use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::mpsc::{Receiver, RecvTimeoutError, Sender, channel}; @@ -9,7 +11,7 @@ use std::time::Duration; use pyo3::exceptions::PyKeyboardInterrupt; use pyo3::prelude::{Bound, PyAny, PyAnyMethods, PyErr, PyResult, Python}; -use pyo3::types::{PyDict, PyDictMethods, PyTuple}; +use pyo3::types::PyBytes; use serde_json::{Map, Value, json}; use sqlbuild_rules::custom_host_pool::constants::{ CGROUP_ROOT, HOST_TIMEOUT_MILLIS, MIN_INVOCATIONS_PER_HOST, PROC_CGROUP, @@ -26,43 +28,18 @@ 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::compiled_project::records::compiled_project_record; use crate::bindings::_helpers::native_host::constants::CUSTOM_HOST_LAUNCH_MODULE; use crate::bindings::_helpers::records::construct::error_class; -use crate::bindings::_helpers::records::construct::new_record; -use crate::bindings::_helpers::records::fields::field as native_field; -use crate::bindings::_helpers::records::fields::replaced_keywords; -use crate::bindings::_helpers::records::ids::{ErrorId, RecordId}; +use crate::bindings::_helpers::records::ids::ErrorId; use crate::bindings::_helpers::rules::constants::{ CUSTOM_HOST_ABANDONED_SUFFIX, CUSTOM_HOST_CANCELLED_MARKER, CUSTOM_HOST_RUNTIME_VERSION, }; use crate::bindings::_helpers::rules::custom_host_observation::{ LaunchLog, ObservedHosts, before_publish, observed_hosts, record_launch, }; +use crate::bindings::_helpers::rules::custom_host_payload::write_payload; 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"; @@ -123,10 +100,10 @@ pub(crate) fn run_custom_hosts( 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 input_path: PathBuf = input_dir.join(format!("{stem}.payload")); let abandoned: PathBuf = input_dir.join(format!("{stem}{}", CUSTOM_HOST_ABANDONED_SUFFIX)); let mut payload: Map = Map::new(); - payload.insert("project_pickle_path".to_owned(), json!(text(&input_path))); + payload.insert("project_payload_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( @@ -142,7 +119,7 @@ pub(crate) fn run_custom_hosts( 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) + write_inputs(&input_path, request) }; let responses: PyResult> = run_partitions( py, @@ -376,19 +353,15 @@ fn host_outcome(spec: &str) -> HostOutcome { 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)?; +/// Publish the project and config payload atomically; hosts started earlier wait for it. +fn write_inputs(path: &Path, request: &HostRequest<'_, '_>) -> PyResult<()> { let parent: &Path = path.parent().unwrap_or(path); let partial: PathBuf = parent.join(format!( - "writing-{}-{}.pickle", + "writing-{}-{}.payload", std::process::id(), INPUT_COUNTER.fetch_add(1, Ordering::Relaxed) )); - let written: PyResult<()> = pickle_into( - &partial, - &PyTuple::new(py, [payload, request.config.clone()])?, - ); + let written: PyResult<()> = write_payload_file(&partial, request); if let Err(error) = written { remove_file(&partial); return Err(error); @@ -397,37 +370,19 @@ fn write_inputs(py: Python<'_>, path: &Path, request: &HostRequest<'_, '_>) -> P Ok(()) } -/// Pickle `value` straight into a new file at `path`, never holding the whole pickle in memory. -fn pickle_into(path: &Path, value: &Bound<'_, PyTuple>) -> PyResult<()> { - let py: Python<'_> = value.py(); - let handle: Bound<'_, PyAny> = imported(py, "builtins", "open")?.call1((text(path), "wb"))?; - let dumped: PyResult> = - imported(py, "pickle", "dump")?.call1((value, &handle)); - let closed: PyResult> = handle.call_method0("close"); - dumped?; - closed?; +/// Stream the payload into a new file at `path`, never holding it whole in memory. +fn write_payload_file(path: &Path, request: &HostRequest<'_, '_>) -> PyResult<()> { + let file: BufWriter = BufWriter::new(File::create(path)?); + let mut file: BufWriter = write_payload(file, request.project, request.config)?; + file.flush()?; 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 mut models: Vec> = Vec::new(); - for model in native_field(project, "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, native_field(&model, name)?)?; - } - models.push(new_record(py, RecordId::CompiledModel, (), Some(&fields))?); - } - let record: Bound<'py, PyAny> = compiled_project_record(project, true)?; - let options: Bound<'py, PyDict> = PyDict::new(py); - options.set_item("models", PyTuple::new(py, models)?)?; - options.set_item("native_project", py.None())?; - replaced_keywords(&record, &options) +/// `pickle.dumps(value)`, for payload values only Python holds. +pub(crate) fn pickled<'py>(value: &Bound<'py, PyAny>) -> PyResult> { + Ok(imported(value.py(), "pickle", "dumps")? + .call1((value,))? + .downcast_into()?) } /// Usable CPUs: the process CPU count capped by any cgroup CPU quota. @@ -489,7 +444,7 @@ fn text(path: &Path) -> String { path.to_string_lossy().into_owned() } -fn rules_error(py: Python<'_>, message: String) -> PyErr { +pub(crate) fn rules_error(py: Python<'_>, message: String) -> PyErr { match error_class(py, ErrorId::Rules).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/fact_replay.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/fact_replay.rs index 28e05abf1b..1fae745e26 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/fact_replay.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/fact_replay.rs @@ -155,12 +155,12 @@ pub(crate) fn is_tree_fact(key: &FactKey) -> bool { pub(crate) struct ModelEntry<'py> { pub(crate) model: Bound<'py, PyAny>, pub(crate) name: String, - path: Bound<'py, PyAny>, + pub(crate) path: Bound<'py, PyAny>, path_text: String, pub(crate) posix: String, pub(crate) config_values: Option>, pub(crate) key: (String, String), - deps: Vec<(String, String)>, + pub(crate) deps: Vec<(String, String)>, } /// The models in project order with the indexes the model facts look them up by. diff --git a/crates/sqlbuild-python/src/bindings/_helpers/rules/mod.rs b/crates/sqlbuild-python/src/bindings/_helpers/rules/mod.rs index 4b2670b1b2..677f355f6e 100644 --- a/crates/sqlbuild-python/src/bindings/_helpers/rules/mod.rs +++ b/crates/sqlbuild-python/src/bindings/_helpers/rules/mod.rs @@ -3,7 +3,9 @@ pub(crate) mod catalogue; pub(crate) mod config_json; pub(crate) mod constants; pub(crate) mod custom_evaluation; +pub(crate) mod custom_host_columns; pub(crate) mod custom_host_observation; +pub(crate) mod custom_host_payload; pub(crate) mod custom_hosts; pub(crate) mod custom_rule_cache; pub(crate) mod custom_rule_evidence; diff --git a/fensu.toml b/fensu.toml index ffd3236e52..d9bfe33c54 100644 --- a/fensu.toml +++ b/fensu.toml @@ -986,7 +986,7 @@ pending_budgets = [ "crates/sqlbuild-python/src/bindings/_helpers/rules/catalogue.rs=7", "crates/sqlbuild-python/src/bindings/_helpers/rules/custom_evaluation.rs=5", "crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_observation.rs=3", - "crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs=14", + "crates/sqlbuild-python/src/bindings/_helpers/rules/custom_hosts.rs=11", "crates/sqlbuild-python/src/bindings/_helpers/rules/custom_rule_cache.rs=1", "crates/sqlbuild-python/src/bindings/_helpers/rules/custom_rule_evidence.rs=21", "crates/sqlbuild-python/src/bindings/_helpers/rules/early_sql_lint.rs=2", diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index f6103ffdbe..9fcce9f310 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -225,6 +225,7 @@ MAX_READ_BYTES: int CURRENT_COMMAND_OUTPUT_SCHEMA_VERSION: int MIGRATION_WRITE_ATTEMPTS: int CUSTOM_HOST_PROTOCOL_VERSION: int +CUSTOM_HOST_PAYLOAD_MAGIC: bytes def plan_sql_tests( *, @@ -6737,7 +6738,8 @@ def evaluate_custom_rules( verify_determinism: bool = False, implementation_fingerprints: dict[str, str] | None = None, ) -> CustomRulesOutcome: ... -def custom_host_payload(project: CompiledProjectView) -> CompiledProject: ... +def custom_host_payload(project: CompiledProjectView, config: object) -> bytes: ... +def custom_host_declared_columns(data: bytes) -> tuple[SchemaColumn, ...]: ... def custom_rule_fact_digest( *, project: CompiledProjectView, project_dir: Path, key: tuple[str, ...] ) -> str | None: ... diff --git a/src/sqlbuild/rule_engine/_helpers/host/custom_evaluation.py b/src/sqlbuild/rule_engine/_helpers/host/custom_evaluation.py index 1345c0b19d..355cf43f44 100644 --- a/src/sqlbuild/rule_engine/_helpers/host/custom_evaluation.py +++ b/src/sqlbuild/rule_engine/_helpers/host/custom_evaluation.py @@ -24,7 +24,6 @@ RulesError, RuleSubject, ) -from sqlbuild.compiler.compile.types import CompiledModelView, CompiledProjectView from sqlbuild.rule_engine._helpers.host.module_state import rule_namespaces from sqlbuild.rule_engine._helpers.host.tracked_facts import tracked_fact_views from sqlbuild.rule_engine.classes.fact_reads import FactReads @@ -37,14 +36,14 @@ ) from sqlbuild.rule_engine.classes.runtime_guard import RuntimeGuard from sqlbuild.rule_engine.constants import CUSTOM_HOST_CANCELLED_MESSAGE -from sqlbuild.rule_engine.types import CustomRulePlan +from sqlbuild.rule_engine.types import CustomRulePlan, RuleFactModel, RuleFactProject type _FindingKey = tuple[str, str, int, int, str, str] @dataclass(frozen=True) class _PassInputs: - project: CompiledProjectView + project: RuleFactProject config: RulesConfig project_dir: Path selected_rules: tuple[Rule, ...] @@ -57,7 +56,7 @@ class _PassInputs: def evaluate_custom_rules( *, - project: CompiledProjectView, + project: RuleFactProject, config: RulesConfig, project_dir: Path, selected_rules: tuple[Rule, ...], @@ -75,7 +74,7 @@ def evaluate_custom_rules( ) if not custom_rules: return CustomRuleRun(evaluations=(), untracked_codes=frozenset()) - models: tuple[CompiledModelView, ...] = tuple( + models: tuple[RuleFactModel, ...] = tuple( sorted(project.models, key=lambda item: item.relative_path.as_posix()) ) facts: RuleFactViews = build_rule_fact_views( @@ -102,7 +101,7 @@ def evaluate_custom_rules( ) def run_pass( - *, rules: tuple[Rule, ...], subjects: tuple[CompiledModelView, ...], active: _ReadTracker + *, rules: tuple[Rule, ...], subjects: tuple[RuleFactModel, ...], active: _ReadTracker ) -> list[CustomRuleEvaluation]: return _evaluate_pass(inputs=inputs, rules=rules, models=subjects, tracker=active) @@ -180,11 +179,11 @@ def _evaluate_pass( *, inputs: _PassInputs, rules: tuple[Rule, ...], - models: tuple[CompiledModelView, ...], + models: tuple[RuleFactModel, ...], tracker: _ReadTracker, ) -> list[CustomRuleEvaluation]: evaluations: list[CustomRuleEvaluation] = [] - project: CompiledProjectView = inputs.project + project: RuleFactProject = inputs.project plan: CustomRulePlan | None = inputs.plan guard: RuntimeGuard | None = inputs.guard context_facts: RuleFactViews = ( diff --git a/src/sqlbuild/rule_engine/_helpers/host/custom_host.py b/src/sqlbuild/rule_engine/_helpers/host/custom_host.py index 7ac2bcc415..fbc81eb9c7 100644 --- a/src/sqlbuild/rule_engine/_helpers/host/custom_host.py +++ b/src/sqlbuild/rule_engine/_helpers/host/custom_host.py @@ -3,22 +3,19 @@ from __future__ import annotations import contextlib -import gc import io import json import os -import pickle import sys import time from dataclasses import asdict from pathlib import Path, PurePath -from typing import Any, BinaryIO +from typing import Any from sqlbuild._native import ( CUSTOM_HOST_LAUNCH_MODULE, CUSTOM_HOST_PROTOCOL_VERSION, CUSTOM_HOST_RUNTIME_VERSION, - CompiledProject, CustomHostPartition, CustomRuleRun, Finding, @@ -28,8 +25,8 @@ RulesError, build_rule_catalogue, ) -from sqlbuild.compiler.compile.types import CompiledProjectView from sqlbuild.rule_engine._helpers.host.custom_evaluation import evaluate_custom_rules +from sqlbuild.rule_engine._helpers.host.host_payload import decode_host_payload from sqlbuild.rule_engine.classes.runtime_guard import RuntimeGuard from sqlbuild.rule_engine.constants import ( CUSTOM_HOST_CANCELLED_EXIT_CODE, @@ -37,11 +34,10 @@ CUSTOM_HOST_EXEC_OS_NAME, CUSTOM_HOST_HASH_SEED, CUSTOM_HOST_INPUT_POLL_SECONDS, - CUSTOM_HOST_INPUT_TUPLE_SIZE, CUSTOM_HOST_INPUT_WAIT_SECONDS, CUSTOM_HOST_MAX_TRACKED_READS, ) -from sqlbuild.rule_engine.types import CustomRulePlan +from sqlbuild.rule_engine.types import CustomRulePlan, RuleFactProject def main() -> int: @@ -121,10 +117,10 @@ def main() -> int: def _decode_inputs( *, payload: dict[str, Any], cancel_marker: Path | None -) -> tuple[CompiledProjectView, RulesConfig]: +) -> tuple[RuleFactProject, RulesConfig]: project_dir: Path = Path(str(payload["project_dir"])).resolve() input_root: Path = (project_dir / "target" / "rules-cache" / "host-inputs").resolve() - input_path: Path = Path(str(payload["project_pickle_path"])) + input_path: Path = Path(str(payload["project_payload_path"])) if not input_path.is_absolute() or not input_path.parent.resolve().is_relative_to(input_root): raise RulesError("custom host project payload path is invalid") await_inputs( @@ -138,16 +134,7 @@ def _decode_inputs( ) if input_path.is_symlink() or not input_path.resolve().is_relative_to(input_root): raise RulesError("custom host project payload path is invalid") - with input_path.open("rb") as handle: - decoded: object = _load_long_lived(handle) - if ( - not isinstance(decoded, tuple) - or len(decoded) != CUSTOM_HOST_INPUT_TUPLE_SIZE - or not isinstance(decoded[0], CompiledProject) - or not isinstance(decoded[1], RulesConfig) - ): - raise RulesError("custom host project payload has invalid types") - return decoded + return decode_host_payload(input_path.read_bytes()) def await_inputs(*, path: Path, cancel_markers: tuple[Path, ...], parent_pid: int) -> None: @@ -187,19 +174,6 @@ def _input_marker(*, payload: dict[str, Any], key: str) -> Path | None: return marker -def _load_long_lived(handle: BinaryIO) -> Any: - """Unpickle the host-lifetime project without cyclic collection, then freeze it.""" - - was_enabled: bool = gc.isenabled() - gc.disable() - try: - return pickle.load(handle) - finally: - gc.freeze() - if was_enabled: - gc.enable() - - def _cancel_marker(payload: dict[str, Any]) -> Path | None: """Return the run's cancellation marker, which must live beside the project payload.""" diff --git a/src/sqlbuild/rule_engine/_helpers/host/host_payload.py b/src/sqlbuild/rule_engine/_helpers/host/host_payload.py new file mode 100644 index 0000000000..c2e4a64844 --- /dev/null +++ b/src/sqlbuild/rule_engine/_helpers/host/host_payload.py @@ -0,0 +1,189 @@ +"""Compiler facts a custom-rule host decodes lazily from its native payload.""" + +from __future__ import annotations + +import json +import pickle +from functools import cached_property +from pathlib import Path +from typing import Any + +from sqlbuild._native import ( + CUSTOM_HOST_PAYLOAD_MAGIC, + CompiledSqlTest, + ConstantDeclaration, + EnumDeclaration, + InferredColumn, + InferredNullability, + RulesConfig, + RulesError, + custom_host_declared_columns, +) +from sqlbuild.compiler.compile.types import ( + CompiledAuditView, + CompiledFunctionView, + CompiledSeedView, + CompiledSourceView, +) +from sqlbuild.rule_engine.types import RuleFactProject + +_TRAILER_SIZE: int = 8 +_PICKLE_TAG: str = "p" +_TEXT_TAG: str = "s" +_TEXTS_TAG: str = "t" +_INFERRED_TAG: str = "r" +_COLUMNS_TAG: str = "c" + + +class _Payload: + """The payload bytes and the values its header names by tag, offset and length.""" + + __slots__ = ("_data",) + + def __init__(self, data: bytes) -> None: + self._data: bytes = data + + def value(self, encoded: Any) -> Any: + if not isinstance(encoded, list): + return encoded + tag: str = encoded[0] + if tag == _TEXTS_TAG: + return tuple(encoded[1]) + start: int = encoded[1] + part: memoryview = memoryview(self._data)[start : start + encoded[2]] + if tag == _PICKLE_TAG: + return pickle.loads(part) + if tag == _TEXT_TAG: + return str(part, "utf-8") + if tag == _INFERRED_TAG: + return tuple( + InferredColumn( + name=name, type=column_type, nullability=InferredNullability(nullability) + ) + for name, column_type, nullability in json.loads(bytes(part)) + ) + if tag == _COLUMNS_TAG: + return custom_host_declared_columns(bytes(part)) + raise RulesError(f"custom host project payload has an unknown part {tag!r}") + + +class _HostModel: + """One compiled model's rule facts, each decoded on first read.""" + + def __init__(self, *, payload: _Payload, entry: dict[str, Any]) -> None: + self._payload: _Payload = payload + self._entry: dict[str, Any] = entry + self._config: dict[str, object] = {} + + @cached_property + def name(self) -> str: + return self._entry["name"] + + @cached_property + def relative_path(self) -> Path: + path: object = self._entry["path"] + return Path(path) if isinstance(path, str) else self._payload.value(path) + + @cached_property + def key(self) -> tuple[str, str]: + resource_type, name = self._entry["key"] + return (resource_type, name) + + @cached_property + def deps(self) -> tuple[tuple[str, str], ...]: + return tuple((resource_type, name) for resource_type, name in self._entry["deps"]) + + @cached_property + def authored_sql(self) -> str: + return self._payload.value(self._entry["authored"]) + + @cached_property + def query_sql(self) -> str: + return self._payload.value(self._entry["query"]) + + def config_value(self, name: str) -> object: + """Return `config.values.get(name)` for one config value rule facts read.""" + + if name not in self._config: + config: dict[str, Any] = self._entry["config"] + self._config[name] = self._payload.value(config[name]) if name in config else None + return self._config[name] + + @cached_property + def declared_columns(self) -> tuple[object, ...]: + encoded: object = self._entry["declared"] + return () if encoded is None else self._payload.value(encoded) + + @cached_property + def inferred_columns(self) -> tuple[InferredColumn, ...] | None: + return self._payload.value(self._entry["inferred"]) + + @cached_property + def enum_declarations(self) -> tuple[EnumDeclaration, ...]: + encoded: object = self._entry["enums"] + return () if encoded is None else self._payload.value(encoded) + + @cached_property + def constant_declarations(self) -> tuple[ConstantDeclaration, ...]: + encoded: object = self._entry["constants"] + return () if encoded is None else self._payload.value(encoded) + + +class _HostProject: + """One compiled project's rule facts, each collection decoded on first read.""" + + def __init__(self, *, payload: _Payload, header: dict[str, Any]) -> None: + self._payload: _Payload = payload + self._header: dict[str, Any] = header + self.models: tuple[_HostModel, ...] = tuple( + _HostModel(payload=payload, entry=entry) for entry in header["models"] + ) + self.effective_target_name: str | None = payload.value(header["target"]) + + def _section(self, name: str) -> Any: + return self._payload.value(self._header["sections"][name]) + + @cached_property + def sources(self) -> tuple[CompiledSourceView, ...]: + return self._section("sources") + + @cached_property + def seeds(self) -> tuple[CompiledSeedView, ...]: + return self._section("seeds") + + @cached_property + def functions(self) -> tuple[CompiledFunctionView, ...]: + return self._section("functions") + + @cached_property + def audits(self) -> tuple[CompiledAuditView, ...]: + return self._section("audits") + + @cached_property + def sql_tests(self) -> tuple[CompiledSqlTest, ...]: + return self._section("sql_tests") + + @cached_property + def public_enums(self) -> dict[str, EnumDeclaration]: + return self._section("public_enums") + + @cached_property + def public_constants(self) -> dict[str, ConstantDeclaration]: + return self._section("public_constants") + + +def decode_host_payload(data: bytes) -> tuple[RuleFactProject, RulesConfig]: + """Read the header of one payload; every fact is decoded only when a rule reads it.""" + + magic: bytes = CUSTOM_HOST_PAYLOAD_MAGIC + if len(data) < len(magic) + _TRAILER_SIZE or not data.startswith(magic): + raise RulesError("custom host project payload is not a payload") + header_start: int = int.from_bytes(data[-_TRAILER_SIZE:], "little") + header: object = json.loads(data[header_start:-_TRAILER_SIZE]) + if not isinstance(header, dict): + raise RulesError("custom host project payload header is invalid") + payload: _Payload = _Payload(data) + config: object = payload.value(header["config"]) + if not isinstance(config, RulesConfig): + raise RulesError("custom host project payload has invalid types") + return _HostProject(payload=payload, header=header), config diff --git a/src/sqlbuild/rule_engine/classes/audit_facts.py b/src/sqlbuild/rule_engine/classes/audit_facts.py index bcbb39994a..8b2c26e559 100644 --- a/src/sqlbuild/rule_engine/classes/audit_facts.py +++ b/src/sqlbuild/rule_engine/classes/audit_facts.py @@ -7,14 +7,18 @@ Model, compiled_audit_record, ) -from sqlbuild.compiler.compile.types import CompiledProjectView +from sqlbuild.rule_engine.types import RuleFactProject class AuditFacts: """Compiled audits indexed by their model subject.""" - def __init__(self, *, project: CompiledProjectView) -> None: - self._audits: tuple[CompiledAudit, ...] = tuple(map(compiled_audit_record, project.audits)) + def __init__(self, *, project: RuleFactProject) -> None: + self._project: RuleFactProject = project + + @cached_property + def _audits(self) -> tuple[CompiledAudit, ...]: + return tuple(map(compiled_audit_record, self._project.audits)) @cached_property def _by_model(self) -> dict[str, tuple[CompiledAudit, ...]]: diff --git a/src/sqlbuild/rule_engine/classes/column_facts.py b/src/sqlbuild/rule_engine/classes/column_facts.py index 6cad2a9fa7..42551fe6c7 100644 --- a/src/sqlbuild/rule_engine/classes/column_facts.py +++ b/src/sqlbuild/rule_engine/classes/column_facts.py @@ -3,22 +3,21 @@ from functools import cached_property from sqlbuild._native import InferredColumn, Model -from sqlbuild.compiler.compile.types import CompiledModelView, CompiledProjectView +from sqlbuild.rule_engine.types import RuleFactModel, RuleFactProject class ColumnFacts: """Declared and compiler-inferred output columns.""" - def __init__(self, *, project: CompiledProjectView) -> None: - self._project: CompiledProjectView = project + def __init__(self, *, project: RuleFactProject) -> None: + self._project: RuleFactProject = project @cached_property - def _models(self) -> dict[object, CompiledModelView]: + def _models(self) -> dict[object, RuleFactModel]: return {model.relative_path: model for model in self._project.models} def declared(self, model: Model) -> tuple[object, ...]: - compiled: CompiledModelView = self._models[model.path] - return () if compiled.schema_entry is None else compiled.schema_entry.columns + return self._models[model.path].declared_columns def inferred(self, model: Model) -> tuple[InferredColumn, ...] | None: return self._models[model.path].inferred_columns diff --git a/src/sqlbuild/rule_engine/classes/contract_facts.py b/src/sqlbuild/rule_engine/classes/contract_facts.py index 454bc611e4..62c46ed76a 100644 --- a/src/sqlbuild/rule_engine/classes/contract_facts.py +++ b/src/sqlbuild/rule_engine/classes/contract_facts.py @@ -1,23 +1,23 @@ """Contract facts for custom rules.""" from sqlbuild._native import CONTRACT_ENFORCED, Model -from sqlbuild.compiler.compile.types import CompiledModelView, CompiledProjectView +from sqlbuild.rule_engine.types import RuleFactModel, RuleFactProject class ContractFacts: """Compiler-resolved model contract facts.""" - def __init__(self, *, project: CompiledProjectView) -> None: - self._models: dict[object, CompiledModelView] = { + def __init__(self, *, project: RuleFactProject) -> None: + self._models: dict[object, RuleFactModel] = { model.relative_path: model for model in project.models } def enforced(self, model: Model) -> bool: - value: object = self._models[model.path].config.values.get("contract") + value: object = self._models[model.path].config_value("contract") return value == CONTRACT_ENFORCED def grain(self, model: Model) -> tuple[str, ...]: - value: object = self._models[model.path].config.values.get("grain") + value: object = self._models[model.path].config_value("grain") if isinstance(value, (tuple, list)) and all(isinstance(item, str) for item in value): return tuple(str(item) for item in value) return () diff --git a/src/sqlbuild/rule_engine/classes/declaration_facts.py b/src/sqlbuild/rule_engine/classes/declaration_facts.py index e0724599f7..0b0ea02421 100644 --- a/src/sqlbuild/rule_engine/classes/declaration_facts.py +++ b/src/sqlbuild/rule_engine/classes/declaration_facts.py @@ -1,21 +1,35 @@ """Declaration facts for custom rules.""" +from functools import cached_property + from sqlbuild._native import ConstantDeclaration, EnumDeclaration -from sqlbuild.compiler.compile.types import CompiledProjectView +from sqlbuild.rule_engine.types import RuleFactProject class DeclarationFacts: """Compiler-resolved public and model-scoped declarations.""" - def __init__(self, *, project: CompiledProjectView) -> None: - self.public_enums: tuple[EnumDeclaration, ...] = tuple(project.public_enums.values()) - self.public_constants: tuple[ConstantDeclaration, ...] = tuple( - project.public_constants.values() - ) + def __init__(self, *, project: RuleFactProject) -> None: + self._project: RuleFactProject = project + + @cached_property + def public_enums(self) -> tuple[EnumDeclaration, ...]: + return tuple(self._project.public_enums.values()) + + @cached_property + def public_constants(self) -> tuple[ConstantDeclaration, ...]: + return tuple(self._project.public_constants.values()) + + @cached_property + def enums(self) -> tuple[EnumDeclaration, ...]: enums: list[EnumDeclaration] = list(self.public_enums) - constants: list[ConstantDeclaration] = list(self.public_constants) - for model in project.models: + for model in self._project.models: enums.extend(model.enum_declarations) + return tuple(enums) + + @cached_property + def constants(self) -> tuple[ConstantDeclaration, ...]: + constants: list[ConstantDeclaration] = list(self.public_constants) + for model in self._project.models: constants.extend(model.constant_declarations) - self.enums: tuple[EnumDeclaration, ...] = tuple(enums) - self.constants: tuple[ConstantDeclaration, ...] = tuple(constants) + return tuple(constants) diff --git a/src/sqlbuild/rule_engine/classes/graph_facts.py b/src/sqlbuild/rule_engine/classes/graph_facts.py index 1ccc5069fd..02538c72a9 100644 --- a/src/sqlbuild/rule_engine/classes/graph_facts.py +++ b/src/sqlbuild/rule_engine/classes/graph_facts.py @@ -3,21 +3,21 @@ from functools import cached_property from sqlbuild._native import Model -from sqlbuild.compiler.compile.types import CompiledModelView, CompiledProjectView from sqlbuild.rule_engine.classes.project_tree import public_model +from sqlbuild.rule_engine.types import RuleFactModel, RuleFactProject class GraphFacts: """Compiler-resolved model dependencies and dependents.""" def __init__( - self, *, project: CompiledProjectView, models: tuple[Model, ...] | None = None + self, *, project: RuleFactProject, models: tuple[Model, ...] | None = None ) -> None: - self._project: CompiledProjectView = project + self._project: RuleFactProject = project self._models: tuple[Model, ...] | None = models @cached_property - def _compiled_by_path(self) -> dict[object, CompiledModelView]: + def _compiled_by_path(self) -> dict[object, RuleFactModel]: return {model.relative_path: model for model in self._project.models} @cached_property @@ -40,9 +40,9 @@ def _downstream(self) -> dict[object, list[Model]]: return downstream def dependencies(self, model: Model) -> tuple[Model, ...]: - compiled: CompiledModelView = self._compiled_by_path[model.path] + compiled: RuleFactModel = self._compiled_by_path[model.path] return tuple(self._by_key[key] for key in compiled.deps if key in self._by_key) def dependents(self, model: Model) -> tuple[Model, ...]: - compiled: CompiledModelView = self._compiled_by_path[model.path] + compiled: RuleFactModel = self._compiled_by_path[model.path] return tuple(sorted(self._downstream.get(compiled.key, ()), key=lambda item: item.name)) diff --git a/src/sqlbuild/rule_engine/classes/model_test_facts.py b/src/sqlbuild/rule_engine/classes/model_test_facts.py index 46efec3754..54878c63c2 100644 --- a/src/sqlbuild/rule_engine/classes/model_test_facts.py +++ b/src/sqlbuild/rule_engine/classes/model_test_facts.py @@ -3,14 +3,19 @@ from functools import cached_property from sqlbuild._native import CompiledSqlTest, Model -from sqlbuild.compiler.compile.types import CompiledProjectView, SqlTestMode +from sqlbuild.compiler.compile.types import SqlTestMode +from sqlbuild.rule_engine.types import RuleFactProject class TestFacts: """Compiled SQL tests indexed by their model subject.""" - def __init__(self, *, project: CompiledProjectView) -> None: - self._tests: tuple[CompiledSqlTest, ...] = project.sql_tests + def __init__(self, *, project: RuleFactProject) -> None: + self._project: RuleFactProject = project + + @cached_property + def _tests(self) -> tuple[CompiledSqlTest, ...]: + return self._project.sql_tests @cached_property def _by_model(self) -> dict[str, tuple[CompiledSqlTest, ...]]: diff --git a/src/sqlbuild/rule_engine/classes/project_facts.py b/src/sqlbuild/rule_engine/classes/project_facts.py index 8c5b37873c..b42cb60b84 100644 --- a/src/sqlbuild/rule_engine/classes/project_facts.py +++ b/src/sqlbuild/rule_engine/classes/project_facts.py @@ -1,5 +1,6 @@ """Whole-project facts for custom rules.""" +from functools import cached_property from pathlib import Path from sqlbuild._native import ( @@ -9,23 +10,39 @@ compiled_seed_record, compiled_source_record, ) -from sqlbuild.compiler.compile.types import CompiledProjectView from sqlbuild.rule_engine.classes.project_tree import ProjectTree, public_model +from sqlbuild.rule_engine.types import RuleFactProject class ProjectFacts: """Compiler-owned project resources and deterministic project tree.""" - def __init__(self, *, project: CompiledProjectView, project_dir: Path) -> None: + def __init__(self, *, project: RuleFactProject, project_dir: Path) -> None: + self._project: RuleFactProject = project self.models: tuple[Model, ...] = tuple( public_model(model) for model in sorted(project.models, key=lambda item: item.relative_path.as_posix()) ) - self.sources: tuple[object, ...] = tuple(map(compiled_source_record, project.sources)) - self.seeds: tuple[object, ...] = tuple(map(compiled_seed_record, project.seeds)) - self.functions: tuple[object, ...] = tuple(map(compiled_function_record, project.functions)) - self.tests: tuple[object, ...] = project.sql_tests - self.audits: tuple[object, ...] = tuple(map(compiled_audit_record, project.audits)) self.tree: ProjectTree = ProjectTree( project_dir=project_dir, project=project, models=self.models ) + + @cached_property + def sources(self) -> tuple[object, ...]: + return tuple(map(compiled_source_record, self._project.sources)) + + @cached_property + def seeds(self) -> tuple[object, ...]: + return tuple(map(compiled_seed_record, self._project.seeds)) + + @cached_property + def functions(self) -> tuple[object, ...]: + return tuple(map(compiled_function_record, self._project.functions)) + + @cached_property + def tests(self) -> tuple[object, ...]: + return self._project.sql_tests + + @cached_property + def audits(self) -> tuple[object, ...]: + return tuple(map(compiled_audit_record, self._project.audits)) diff --git a/src/sqlbuild/rule_engine/classes/project_tree.py b/src/sqlbuild/rule_engine/classes/project_tree.py index c2b62eedf9..9fc3d144fa 100644 --- a/src/sqlbuild/rule_engine/classes/project_tree.py +++ b/src/sqlbuild/rule_engine/classes/project_tree.py @@ -3,6 +3,7 @@ from __future__ import annotations from pathlib import Path +from typing import cast import sqlbuild._native as _native from sqlbuild._native import ( @@ -16,7 +17,7 @@ project_tree_read_text, project_tree_relative_parts, ) -from sqlbuild.compiler.compile.types import CompiledProjectView +from sqlbuild.rule_engine.types import RuleFactModel, RuleFactProject class ProjectTree: @@ -26,7 +27,7 @@ def __init__( self, *, project_dir: Path, - project: CompiledProjectView, + project: RuleFactProject, models: tuple[Model, ...] | None = None, ) -> None: self._root: Path = project_dir.resolve() @@ -74,11 +75,20 @@ def read_text(self, path: str) -> str: def public_model(model: object) -> Model: from sqlbuild._native import CompiledModel - if not isinstance(model, (CompiledModel, _native.NativeCompiledModel)): - raise RuleUsageError("model is not a compiled SQL model") - value: object = model.config.values.get("materialized") - return Model( - name=model.name, - path=model.relative_path, - materialization=value if isinstance(value, str) else None, - ) + if isinstance(model, (CompiledModel, _native.NativeCompiledModel)): + value: object = model.config.values.get("materialized") + return Model( + name=model.name, + path=model.relative_path, + materialization=value if isinstance(value, str) else None, + ) + facts: RuleFactModel = cast(RuleFactModel, model) + try: + value = facts.config_value("materialized") + return Model( + name=facts.name, + path=facts.relative_path, + materialization=value if isinstance(value, str) else None, + ) + except AttributeError as error: + raise RuleUsageError("model is not a compiled SQL model") from error diff --git a/src/sqlbuild/rule_engine/classes/rule_context.py b/src/sqlbuild/rule_engine/classes/rule_context.py index 555d7db93e..5700d2614b 100644 --- a/src/sqlbuild/rule_engine/classes/rule_context.py +++ b/src/sqlbuild/rule_engine/classes/rule_context.py @@ -17,7 +17,6 @@ SqlNode, finding_position, ) -from sqlbuild.compiler.compile.types import CompiledModelView, CompiledProjectView from sqlbuild.rule_engine.classes.audit_facts import AuditFacts from sqlbuild.rule_engine.classes.column_facts import ColumnFacts from sqlbuild.rule_engine.classes.contract_facts import ContractFacts @@ -27,10 +26,11 @@ from sqlbuild.rule_engine.classes.model_test_facts import TestFacts from sqlbuild.rule_engine.classes.project_facts import ProjectFacts from sqlbuild.rule_engine.classes.sql_facts import SqlFacts +from sqlbuild.rule_engine.types import RuleFactModel, RuleFactProject def build_rule_fact_views( - *, project: CompiledProjectView, project_dir: Path, dialect: str + *, project: RuleFactProject, project_dir: Path, dialect: str ) -> RuleFactViews: """Build immutable views once for one compiled project snapshot.""" @@ -53,10 +53,10 @@ class EvaluationRuleContext: def __init__( self, *, - model: CompiledModelView | None, + model: RuleFactModel | None, rule: Rule, config: RulesConfig, - project: CompiledProjectView, + project: RuleFactProject, project_dir: Path, selected_rules: tuple[Rule, ...], dialect: str, diff --git a/src/sqlbuild/rule_engine/classes/sql_facts.py b/src/sqlbuild/rule_engine/classes/sql_facts.py index c600b94107..8a3e994583 100644 --- a/src/sqlbuild/rule_engine/classes/sql_facts.py +++ b/src/sqlbuild/rule_engine/classes/sql_facts.py @@ -3,27 +3,27 @@ from functools import cached_property from sqlbuild._native import Model, ModelSql, RuleUsageError -from sqlbuild.compiler.compile.types import CompiledModelView, CompiledProjectView from sqlbuild.rule_engine.classes.sql_document import SqlDocument +from sqlbuild.rule_engine.types import RuleFactModel, RuleFactProject class SqlFacts: """Authored and expanded SQL indexed by public model identity.""" - def __init__(self, *, dialect: str, project: CompiledProjectView) -> None: + def __init__(self, *, dialect: str, project: RuleFactProject) -> None: self._dialect: str = dialect - self._project: CompiledProjectView = project + self._project: RuleFactProject = project self._documents: dict[object, ModelSql] = {} @cached_property - def _models(self) -> dict[object, CompiledModelView]: + def _models(self) -> dict[object, RuleFactModel]: return {model.relative_path: model for model in self._project.models} def for_model(self, model: Model) -> ModelSql: cached: ModelSql | None = self._documents.get(model.path) if cached is not None: return cached - compiled: CompiledModelView = self._resolve(model) + compiled: RuleFactModel = self._resolve(model) documents: ModelSql = ModelSql( authored=SqlDocument(source=compiled.authored_sql, dialect=self._dialect), expanded=SqlDocument(source=compiled.query_sql, dialect=self._dialect), @@ -31,7 +31,7 @@ def for_model(self, model: Model) -> ModelSql: self._documents[model.path] = documents return documents - def _resolve(self, model: Model) -> CompiledModelView: + def _resolve(self, model: Model) -> RuleFactModel: try: return self._models[model.path] except KeyError as error: diff --git a/src/sqlbuild/rule_engine/constants.py b/src/sqlbuild/rule_engine/constants.py index 02533e5f14..e7275bc033 100644 --- a/src/sqlbuild/rule_engine/constants.py +++ b/src/sqlbuild/rule_engine/constants.py @@ -1,7 +1,6 @@ """Rules engine constants.""" CUSTOM_HOST_MAX_TRACKED_READS: int = 100_000 -CUSTOM_HOST_INPUT_TUPLE_SIZE: int = 2 CUSTOM_HOST_INPUT_POLL_SECONDS: float = 0.002 CUSTOM_HOST_INPUT_WAIT_SECONDS: float = 120.0 CUSTOM_HOST_MODULE: str = "sqlbuild.rule_engine._helpers.host.custom_host" diff --git a/src/sqlbuild/rule_engine/types.py b/src/sqlbuild/rule_engine/types.py index fec5a01fac..5f052e4ae9 100644 --- a/src/sqlbuild/rule_engine/types.py +++ b/src/sqlbuild/rule_engine/types.py @@ -7,7 +7,23 @@ from typing import TYPE_CHECKING, Protocol if TYPE_CHECKING: - from sqlbuild._native import Finding, Model, ProjectPath, RuleOption, SqlNode + from sqlbuild._native import ( + CompiledSqlTest, + ConstantDeclaration, + EnumDeclaration, + Finding, + InferredColumn, + Model, + ProjectPath, + RuleOption, + SqlNode, + ) + from sqlbuild.compiler.compile.types import ( + CompiledAuditView, + CompiledFunctionView, + CompiledSeedView, + CompiledSourceView, + ) from sqlbuild.rule_engine.classes.audit_facts import AuditFacts from sqlbuild.rule_engine.classes.column_facts import ColumnFacts from sqlbuild.rule_engine.classes.contract_facts import ContractFacts @@ -22,6 +38,73 @@ type CustomRulePlan = dict[str, frozenset[str] | None] +class RuleFactModel(Protocol): + """One compiled model's facts as a custom-rule host reads them.""" + + @property + def name(self) -> str: ... + + @property + def relative_path(self) -> Path: ... + + @property + def key(self) -> object: ... + + @property + def deps(self) -> tuple[object, ...]: ... + + @property + def authored_sql(self) -> str: ... + + @property + def query_sql(self) -> str: ... + + @property + def declared_columns(self) -> tuple[object, ...]: ... + + @property + def inferred_columns(self) -> tuple[InferredColumn, ...] | None: ... + + @property + def enum_declarations(self) -> tuple[EnumDeclaration, ...]: ... + + @property + def constant_declarations(self) -> tuple[ConstantDeclaration, ...]: ... + + def config_value(self, name: str) -> object: ... + + +class RuleFactProject(Protocol): + """One compiled project's facts as a custom-rule host reads them.""" + + @property + def models(self) -> tuple[RuleFactModel, ...]: ... + + @property + def effective_target_name(self) -> str | None: ... + + @property + def sources(self) -> tuple[CompiledSourceView, ...]: ... + + @property + def seeds(self) -> tuple[CompiledSeedView, ...]: ... + + @property + def functions(self) -> tuple[CompiledFunctionView, ...]: ... + + @property + def audits(self) -> tuple[CompiledAuditView, ...]: ... + + @property + def sql_tests(self) -> tuple[CompiledSqlTest, ...]: ... + + @property + def public_enums(self) -> dict[str, EnumDeclaration]: ... + + @property + def public_constants(self) -> dict[str, ConstantDeclaration]: ... + + class RuleContext(Protocol): """Stable typed compiler-fact views available to custom rules.""" 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 28b083bba4..6dc1cc9791 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py @@ -1177,3 +1177,24 @@ class TemplatedHeaderDiagnosticsTestCase: amount_type: str amount_description: str expected_location: dict[str, object] + + +@dataclass(frozen=True) +class CustomRuleSurfaceTestCase: + """Compiles after each `{path: contents}` edit with rules reading every public fact.""" + + description: str + project_files: dict[str, str] + edits: tuple[dict[str, str], ...] + expected_messages: tuple[str, ...] + expected_hits_and_misses: tuple[tuple[int, int], ...] + + +@dataclass(frozen=True) +class CustomHostOutcomeTestCase: + """A compile whose custom rule reads facts, then fails, exits its host, or finishes slowly.""" + + description: str + rule_body: str + expected_exit_code: int + expected_messages: tuple[str, ...] 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 2ea7a60274..e8b679ba72 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py @@ -25,8 +25,10 @@ from tests.integration.src.sqlbuild.cli.commands.main._test_types import ( CustomHostFailureIntegrationTestCase, CustomHostInterruptIntegrationTestCase, + CustomHostOutcomeTestCase, CustomHostSplitIntegrationTestCase, CustomRuleCacheSequenceTestCase, + CustomRuleSurfaceTestCase, DynamicPivotRulesIntegrationTestCase, EarlyLintInputsTestCase, EmptyRuleSelectionTestCase, @@ -74,6 +76,112 @@ def final_directory(*, model: Model, ctx: RuleContext) -> list[Finding]: _CUSTOM_CACHE_MODEL_SQL: str = 'MODEL (description "Orders");\nSELECT 1 AS order_id\n' _CUSTOM_CACHE_EDITED_SQL: str = _CUSTOM_CACHE_MODEL_SQL.replace("SELECT 1", "SELECT 2") +_SURFACE_RULES: str = """from sqlbuild.rules import Finding, Model, Project, RuleContext, rule + + +def _names(models: tuple[Model, ...]) -> str: + return ",".join(model.name for model in models) + + +@rule(code="XSQBRSURF001", message="model facts", remediation="Read every model fact.") +def model_facts(*, model: Model, ctx: RuleContext) -> list[Finding]: + sql = ctx.sql.for_model(model) + declared = ";".join( + f"{column.name}:{column.type}:{column.nullable}:{column.description}:" + + "/".join( + f"{audit.definition_name}:{sorted(audit.arguments.items())}:{audit.severity}" + for audit in column.audits + ) + + f":{column.location.path.as_posix()}:{column.location.line}" + for column in ctx.columns.declared(model) + ) + inferred = ctx.columns.inferred(model) or () + parts = ( + model.name, + model.path.as_posix(), + str(model.materialization), + str(len(sql.authored.source)), + str(len(sql.expanded.source)), + _names(ctx.graph.dependencies(model)), + _names(ctx.graph.dependents(model)), + declared, + ";".join(f"{c.name}:{c.type}:{c.nullability.value}" for c in inferred), + ",".join(ctx.columns.names(model) or ()), + str(ctx.contracts.enforced(model)), + ",".join(ctx.contracts.grain(model)), + ",".join(test.name for test in ctx.tests.for_model(model)), + ",".join(audit.name for audit in ctx.audits.for_model(model)), + ) + return [ctx.finding(subject=model, message="|".join(parts))] + + +@rule(code="XSQBRSURF002", message="project facts", remediation="Read every project fact.") +def project_facts(*, project: Project, ctx: RuleContext) -> list[Finding]: + facts = ctx.project + parts = ( + str(project.model_count), + str(project.target_name), + _names(facts.models), + ",".join(source.name for source in facts.sources), + ",".join(seed.name for seed in facts.seeds), + ",".join(function.name for function in facts.functions), + ",".join(test.name for test in facts.tests), + ",".join(audit.name for audit in facts.audits), + str(len(ctx.tests.all())), + str(len(ctx.audits.all())), + ",".join(item.name for item in ctx.declarations.public_enums), + ",".join(item.name for item in ctx.declarations.public_constants), + str(len(ctx.declarations.enums)), + str(len(ctx.declarations.constants)), + _names(facts.tree.resources_under("models/marts")), + ) + return [ctx.finding(subject=facts.models[0], message="|".join(parts))] +""" +_SURFACE_FILES: dict[str, str] = { + "sqlbuild_project.toml": ( + 'name = "orders"\nadapter = "duckdb"\n\n[rules]\nselect = ["XSQBRSURF"]\n\n' + "[rules.thresholds]\nmin_custom_rule_test_cases = 0\n" + ), + "rules/surface.py": _SURFACE_RULES, + "sources/raw.yml": ( + "sources:\n - name: raw_orders\n description: Raw orders.\n" + ' expression: "(SELECT 1 AS id, CAST(10.5 AS DOUBLE) AS amount)"\n' + ), + "seeds/regions.csv": "region_id,name\n1,north\n", + "seeds/regions.yml": ( + "seeds:\n - name: regions\n description: Regions.\n columns:\n" + " - name: region_id\n type: INTEGER\n - name: name\n type: VARCHAR\n" + ), + "models/staging/_sqlbuild/_constants/policy.sql": ( + "CONSTANT (name minimum_amount, value 7);\n" + ), + "models/staging/_sqlbuild/_enums/order_status.sql": ( + "ENUM (name order_status, members [PAID, OPEN]);\n" + ), + "models/staging/stg_orders.sql": ( + "MODEL (description 'Staged orders.', contract enforced, columns (\n" + " order_id (type INTEGER, nullable false, description 'Order key',\n" + " audits [not_null, accepted_values (values [1, 2], severity warn)]),\n" + " amount (type DOUBLE),\n status (type order_status),\n));\n\n" + 'SELECT CAST(id AS INTEGER) AS order_id, amount, @enum("order_status").PAID AS status\n' + 'FROM __source("raw_orders") WHERE amount > @const("minimum_amount")\n' + ), + "models/marts/order_totals.sql": ( + 'MODEL (description "Order totals", materialized table);\n\n' + 'SELECT o.order_id, o.amount * 2 AS doubled_amount FROM __ref("stg_orders") AS o\n' + ), + "tests/unit/test_order_totals.sql": ( + "TEST ();\n\nWITH\n__ref__stg_orders AS (SELECT 1 AS order_id, 10.0 AS amount),\n" + "__expected__order_totals AS (SELECT 1 AS order_id, 20.0 AS doubled_amount)\n" + "SELECT 1\n" + ), +} +_SURFACE_STAGING_EDIT: dict[str, str] = { + "models/staging/stg_orders.sql": _SURFACE_FILES["models/staging/stg_orders.sql"].replace( + "AS order_id, amount,", "AS order_id,\n amount," + ) +} + _EARLY_LINT_SLOW_HOOKS: int = 8 _EARLY_LINT_SLOW_EXPANSION_S: float = 0.5 _EARLY_LINT_EXPANSION_DELAYS_S: dict[str, float] = {"orders": 0.0} @@ -1294,7 +1402,7 @@ def final_directory(*, model: Model, ctx: RuleContext) -> list[Finding]: assert warm["compile_timings"]["rule_cache_hits"] >= 1 assert not (tmp_path / "target" / "compiled").exists() host_inputs: Path = tmp_path / "target" / "rules-cache" / "host-inputs" - assert not tuple(host_inputs.glob("*.pickle")) + assert not tuple(host_inputs.glob("*.payload")) @pytest.mark.parametrize( @@ -3006,3 +3114,123 @@ def test_given_custom_rule_cache_when_inputs_change_then_only_unchanged_runs_hit observed.append((timings["rule_cache_hits"], timings["rule_cache_misses"])) assert tuple(observed) == test_case.expected_hits_and_misses + + +@pytest.mark.parametrize( + "test_case", + [ + CustomRuleSurfaceTestCase( + description=( + "rules reading every public fact see the compiled values cold, reuse them warm, " + "and rerun only the edited model's invocation" + ), + project_files=_SURFACE_FILES, + edits=({}, {}, _SURFACE_STAGING_EDIT), + expected_messages=( + "2|None|order_totals,stg_orders|raw_orders|regions||test_order_totals|" + "not_null,accepted_values,accepted_values|1|3|order_status|minimum_amount|1|1|" + "order_totals", + "order_totals|models/marts/order_totals.sql|table|137|79|stg_orders|||" + "order_id:INT:non_null;doubled_amount:DOUBLE:unknown|order_id,doubled_amount|" + "False||test_order_totals|", + "stg_orders|models/staging/stg_orders.sql|None|417|109||order_totals|" + "order_id:INTEGER:False:Order key:not_null:[]:None/accepted_values:" + "[('values', [1, 2])]:warn:models/staging/stg_orders.sql:2;" + "amount:DOUBLE:None:None::models/staging/stg_orders.sql:4;" + "status:VARCHAR:None:None:accepted_values:[('values', ('PAID', 'OPEN'))]:None:" + "models/staging/stg_orders.sql:5|" + "order_id:INT:unknown;amount:DOUBLE:unknown;status:TEXT:non_null|" + "order_id,amount,status|True|||not_null,accepted_values,accepted_values", + ), + expected_hits_and_misses=((0, 4), (4, 0), (3, 1)), + ) + ], + ids=lambda case: case.description, +) +def test_given_rules_reading_every_fact_when_compiling_then_values_and_cache_reuse_hold( + test_case: CustomRuleSurfaceTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + write_project_files(project_dir=tmp_path, files=test_case.project_files) + observed: list[tuple[int, int]] = [] + messages: list[tuple[str, ...]] = [] + for edit in test_case.edits: + write_project_files(project_dir=tmp_path, files=edit) + _ = main(["--project-dir", str(tmp_path), "compile", "--json"]) + output: dict[str, Any] = json.loads(capsys.readouterr().out) + timings: dict[str, int] = output["compile_timings"] + observed.append((timings["rule_cache_hits"], timings["rule_cache_misses"])) + messages.append(tuple(sorted(item["message"] for item in output["diagnostics"]))) + config_path: Path = tmp_path / "sqlbuild_project.toml" + config_path.write_text( + config_path.read_text(encoding="utf-8") + "\n[rules.cache]\nenabled = false\n", + encoding="utf-8", + ) + _ = main(["--project-dir", str(tmp_path), "compile", "--json", "--no-cache"]) + uncached: dict[str, Any] = json.loads(capsys.readouterr().out) + + assert tuple(observed) == test_case.expected_hits_and_misses + assert messages[0] == messages[1] == test_case.expected_messages + assert messages[2] == tuple(sorted(item["message"] for item in uncached["diagnostics"])) + + +@pytest.mark.parametrize( + "test_case", + [ + CustomHostOutcomeTestCase( + description="a rule failing after lazy fact reads reports its model", + rule_body=( + "columns = ctx.columns.inferred(model)\n" + 'if model.name == "orders_001":\n' + ' raise ValueError(f"cannot inspect {columns[0].name}")\n' + "return []" + ), + expected_exit_code=1, + expected_messages=( + "invalid rules lifecycle: custom host failed: rule XSQBRHOST001 failed for " + "models/orders_001.sql: cannot inspect order_id", + ), + ), + CustomHostOutcomeTestCase( + description="a rule exiting its host reports the host failure", + rule_body=( + "sql = ctx.sql.for_model(model).expanded.source\n" + 'if model.name == "orders_002":\n' + " raise SystemExit(3)\n" + "return []" + ), + expected_exit_code=1, + expected_messages=( + "invalid rules lifecycle: custom host failed: process exited with exit status: 3", + ), + ), + CustomHostOutcomeTestCase( + description="a slow rule reading lazy facts reports every finding", + rule_body=( + "_ = sum(range(2_000_000))\n" + "return [ctx.finding(subject=model, message=ctx.sql.for_model(model).expanded.source)]" + ), + expected_exit_code=1, + expected_messages=( + "SELECT 0 AS order_id", + "SELECT 1 AS order_id", + "SELECT 2 AS order_id", + ), + ), + ], + ids=lambda case: case.description, +) +def test_given_custom_rule_outcome_when_compiling_then_result_and_cleanup_are_reported( + test_case: CustomHostOutcomeTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + write_custom_host_project(project_dir=tmp_path, model_count=3, rule_body=test_case.rule_body) + + exit_code: int = main(["--project-dir", str(tmp_path), "compile", "--json", "--no-cache"]) + output: dict[str, Any] = json.loads(capsys.readouterr().out) + + assert exit_code == test_case.expected_exit_code + assert tuple(item["message"] for item in output["diagnostics"]) == test_case.expected_messages + assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) diff --git a/tests/integration/src/sqlbuild/rule_engine/host_payload/_test_types.py b/tests/integration/src/sqlbuild/rule_engine/host_payload/_test_types.py index f3e7a28947..ef6ca263af 100644 --- a/tests/integration/src/sqlbuild/rule_engine/host_payload/_test_types.py +++ b/tests/integration/src/sqlbuild/rule_engine/host_payload/_test_types.py @@ -1,9 +1,25 @@ +from collections.abc import Callable from dataclasses import dataclass +from sqlbuild.compiler.compile.types import CompiledProjectView + @dataclass(frozen=True) class HostPayloadFactsTestCase: + """A compiled project written as a host payload, and the facts that must not be empty.""" + + description: str + project_files: dict[str, str] + payload_source: Callable[[CompiledProjectView], CompiledProjectView] + expected_populated_facts: tuple[str, ...] + + +@dataclass(frozen=True) +class HostPayloadLazyReadTestCase: + """A rule reading one fact per model, and the facts its host must leave encoded.""" + description: str project_files: dict[str, str] - expected_dropped_project_fields: frozenset[str] - expected_dropped_model_fields: frozenset[str] + expected_enforced: tuple[bool, ...] + expected_encoded_model_facts: frozenset[str] + expected_encoded_project_facts: frozenset[str] diff --git a/tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py b/tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py index a4d204b5da..1d6805bf08 100644 --- a/tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py +++ b/tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py @@ -1,84 +1,19 @@ -import dataclasses -from itertools import product +from dataclasses import replace from pathlib import Path +from typing import Any from sqlbuild._native import ( Model, + RuleFactViews, build_compiled_project, - custom_rule_fact_digest, + compiled_model_record, + compiled_project_record, discover_project_inputs, ) from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter -from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.rule_engine.classes.project_tree import public_model -from sqlbuild.rule_engine.constants import ( - FACT_AUDITS_ALL, - FACT_AUDITS_FOR_MODEL, - FACT_COLUMNS_DECLARED, - FACT_COLUMNS_INFERRED, - FACT_COLUMNS_NAMES, - FACT_CONTRACTS_ENFORCED, - FACT_CONTRACTS_GRAIN, - FACT_DECLARATIONS_CONSTANTS, - FACT_DECLARATIONS_ENUMS, - FACT_DECLARATIONS_PUBLIC_CONSTANTS, - FACT_DECLARATIONS_PUBLIC_ENUMS, - FACT_GRAPH_DEPENDENCIES, - FACT_GRAPH_DEPENDENTS, - FACT_PROJECT_AUDITS, - FACT_PROJECT_FUNCTIONS, - FACT_PROJECT_MODELS, - FACT_PROJECT_SEEDS, - FACT_PROJECT_SOURCES, - FACT_PROJECT_TESTS, - FACT_SQL_FOR_MODEL, - FACT_TESTS_ALL, - FACT_TESTS_FOR_MODEL, - FACT_TREE_CHILDREN, - FACT_TREE_DESCENDANTS, - FACT_TREE_GLOB, - FACT_TREE_PATHS, - FACT_TREE_READ_TEXT, - FACT_TREE_RELATIVE_PARTS, - FACT_TREE_RESOURCES_UNDER, -) -from sqlbuild.rule_engine.types import FactKey - -_TREE_ARGUMENTS: tuple[str, ...] = ("models/marts/order_totals.sql", "models") -_GLOBAL_FACTS: tuple[str, ...] = ( - FACT_PROJECT_MODELS, - FACT_PROJECT_SOURCES, - FACT_PROJECT_SEEDS, - FACT_PROJECT_FUNCTIONS, - FACT_PROJECT_TESTS, - FACT_PROJECT_AUDITS, - FACT_TREE_PATHS, - FACT_TESTS_ALL, - FACT_AUDITS_ALL, - FACT_DECLARATIONS_PUBLIC_ENUMS, - FACT_DECLARATIONS_PUBLIC_CONSTANTS, - FACT_DECLARATIONS_ENUMS, - FACT_DECLARATIONS_CONSTANTS, -) -_MODEL_PATH_FACTS: tuple[str, ...] = ( - FACT_SQL_FOR_MODEL, - FACT_GRAPH_DEPENDENCIES, - FACT_GRAPH_DEPENDENTS, - FACT_COLUMNS_DECLARED, - FACT_COLUMNS_INFERRED, - FACT_COLUMNS_NAMES, - FACT_CONTRACTS_ENFORCED, - FACT_CONTRACTS_GRAIN, -) -_MODEL_NAME_FACTS: tuple[str, ...] = (FACT_TESTS_FOR_MODEL, FACT_AUDITS_FOR_MODEL) -_TEXT_FACTS: tuple[tuple[str, int], ...] = ( - (FACT_TREE_CHILDREN, 1), - (FACT_TREE_DESCENDANTS, 1), - (FACT_TREE_GLOB, 1), - (FACT_TREE_RELATIVE_PARTS, 2), - (FACT_TREE_RESOURCES_UNDER, 1), - (FACT_TREE_READ_TEXT, 1), -) +from sqlbuild.compiler.compile.types import CompiledModelView, CompiledProjectView +from sqlbuild.rule_engine.classes.rule_context import build_rule_fact_views +from sqlbuild.rule_engine.types import RuleFactProject class OrdersReferenceResolver: @@ -113,6 +48,41 @@ def resolve_reference( return None +class ReferenceModel: + """A compiled model read the way rule facts read full `CompiledModel` records.""" + + def __init__(self, model: CompiledModelView) -> None: + self._model: CompiledModelView = model + self.name: str = model.name + self.relative_path: Path = model.relative_path + self.key: object = model.key + self.deps: tuple[object, ...] = model.deps + self.authored_sql: str = model.authored_sql + self.query_sql: str = model.query_sql + self.inferred_columns: object = model.inferred_columns + self.enum_declarations: tuple[object, ...] = model.enum_declarations + self.constant_declarations: tuple[object, ...] = model.constant_declarations + self.declared_columns: tuple[object, ...] = getattr(model.schema_entry, "columns", ()) + + def config_value(self, name: str) -> object: + return self._model.config.values.get(name) + + +class ReferenceProject: + """A compiled project read the way rule facts read a full `CompiledProject` record.""" + + def __init__(self, project: CompiledProjectView) -> None: + self.models: tuple[ReferenceModel, ...] = tuple(map(ReferenceModel, project.models)) + self.effective_target_name: str | None = project.effective_target_name + self.sources: tuple[object, ...] = project.sources + self.seeds: tuple[object, ...] = project.seeds + self.functions: tuple[object, ...] = project.functions + self.audits: tuple[object, ...] = project.audits + self.sql_tests: tuple[object, ...] = project.sql_tests + self.public_enums: dict[str, object] = dict(project.public_enums) + self.public_constants: dict[str, object] = dict(project.public_constants) + + def compile_files(*, root: Path, files: dict[str, str]) -> CompiledProjectView: """Write and compile one DuckDB project with an external reference resolver.""" @@ -127,29 +97,55 @@ def compile_files(*, root: Path, files: dict[str, str]) -> CompiledProjectView: ) -def dropped_fields(*, before: object, after: object) -> frozenset[str]: - """Return the fields of `after`'s dataclass whose values differ between the two snapshots.""" +def fact_views(*, project: RuleFactProject, root: Path) -> RuleFactViews: + """Build the fact views one host evaluation reads through.""" - names: tuple[str, ...] = tuple(field.name for field in dataclasses.fields(type(after))) - return frozenset(filter(lambda name: getattr(before, name) != getattr(after, name), names)) + return build_rule_fact_views(project=project, project_dir=root, dialect="duckdb") -def every_fact_key(project: CompiledProjectView) -> tuple[FactKey, ...]: - """Return a key for every fact Rules can read, over every model of the project.""" +def native_project(project: CompiledProjectView) -> CompiledProjectView: + """The compiled project as the compiler built it.""" - models: tuple[Model, ...] = tuple(public_model(model) for model in project.models) - return ( - *((fact,) for fact in _GLOBAL_FACTS), - *((fact, model.path.as_posix()) for fact, model in product(_MODEL_PATH_FACTS, models)), - *((fact, model.name) for fact, model in product(_MODEL_NAME_FACTS, models)), - *((fact, *_TREE_ARGUMENTS[:arity]) for fact, arity in _TEXT_FACTS), - ) + return project + + +def record_project(project: CompiledProjectView) -> CompiledProjectView: + """The compiled project as a `CompiledProject` record holding `CompiledModel` records.""" + + record: Any = compiled_project_record(project) + return replace(record, models=tuple(map(compiled_model_record, record.models))) -def fact_digests(*, project: CompiledProjectView, root: Path) -> dict[FactKey, str | None]: - """Digest every fact the custom-rule cache can key a result on.""" +def every_fact_value(views: RuleFactViews) -> dict[str, object]: + """Read every public fact once, keyed by fact and subject.""" - return { - key: custom_rule_fact_digest(project=project, project_dir=root, key=key) - for key in every_fact_key(project) + values: dict[str, object] = { + "project.models": views.project.models, + "project.sources": views.project.sources, + "project.seeds": views.project.seeds, + "project.functions": views.project.functions, + "project.tests": views.project.tests, + "project.audits": views.project.audits, + "tests.all": views.tests.all(), + "audits.all": views.audits.all(), + "declarations.public_enums": views.declarations.public_enums, + "declarations.public_constants": views.declarations.public_constants, + "declarations.enums": views.declarations.enums, + "declarations.constants": views.declarations.constants, + "tree.resources_under": views.project.tree.resources_under("models"), } + model: Model + for model in views.project.models: + subject: str = model.path.as_posix() + documents: Any = views.sql.for_model(model) + values[f"sql.for_model:{subject}"] = (documents.authored.source, documents.expanded.source) + values[f"graph.dependencies:{subject}"] = views.graph.dependencies(model) + values[f"graph.dependents:{subject}"] = views.graph.dependents(model) + values[f"columns.declared:{subject}"] = views.columns.declared(model) + values[f"columns.inferred:{subject}"] = views.columns.inferred(model) + values[f"columns.names:{subject}"] = views.columns.names(model) + values[f"contracts.enforced:{subject}"] = views.contracts.enforced(model) + values[f"contracts.grain:{subject}"] = views.contracts.grain(model) + values[f"tests.for_model:{subject}"] = views.tests.for_model(model) + values[f"audits.for_model:{subject}"] = views.audits.for_model(model) + return values 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 6d50b35423..29b6d4b5b6 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 @@ -1,23 +1,96 @@ -"""Every Rule fact reads the same values from the trimmed host payload as from the full project.""" +"""Every Rule fact reads the same values from the native host payload as from the full project.""" from pathlib import Path import pytest -from sqlbuild._native import ( - compiled_project_record, - custom_host_payload, -) +from sqlbuild._native import RuleFactViews, RulesConfig, custom_host_payload from sqlbuild.compiler.compile.types import CompiledProjectView -from sqlbuild.rule_engine.classes.project_tree import public_model -from sqlbuild.rule_engine.types import FactKey +from sqlbuild.rule_engine._helpers.host.host_payload import decode_host_payload +from sqlbuild.rule_engine.types import RuleFactProject from tests.integration.src.sqlbuild.rule_engine.host_payload._test_types import ( HostPayloadFactsTestCase, + HostPayloadLazyReadTestCase, ) from tests.integration.src.sqlbuild.rule_engine.host_payload.helpers import ( + ReferenceProject, compile_files, - dropped_fields, - fact_digests, + every_fact_value, + fact_views, + native_project, + record_project, +) + +_PROJECT_FILES: dict[str, str] = { + "sqlbuild_project.toml": 'name = "orders"\nadapter = "duckdb"\n', + "sources/raw.yml": ( + "sources:\n" + " - name: raw_orders\n description: Raw orders.\n" + ' expression: "(SELECT 1 AS id, CAST(10.5 AS DOUBLE) AS amount)"\n' + " columns:\n - name: id\n type: INTEGER\n" + " - name: raw_customers\n description: Raw customers.\n" + ' expression: "(SELECT 1 AS customer_id)"\n' + ), + "seeds/regions.csv": "region_id,name\n1,north\n", + "seeds/regions.yml": ( + "seeds:\n - name: regions\n description: Regions.\n columns:\n" + " - name: region_id\n type: INTEGER\n" + " - name: name\n type: VARCHAR\n" + ), + "functions/sql/is_paid.sql": ( + "FUNCTION (\n description 'Paid status.',\n arguments (status STRING),\n" + " returns BOOLEAN,\n);\n\nstatus = 'PAID'\n" + ), + "constants/policy.sql": "CONSTANT (name minimum_amount, value 7);\n", + "enums/order_status.sql": "ENUM (name order_status, members [PAID, OPEN]);\n", + "models/marts/_sqlbuild/_macros/amounts.py": ( + 'def doubled(amount: str) -> str:\n return f"({amount}) * 2"\n' + ), + "models/staging/stg_orders.sql": ( + "MODEL (description 'Staged orders.',\n contract enforced,\n" + " columns (\n" + " order_id (type INTEGER, nullable false, description 'Order key',\n" + " audits [not_null, accepted_values (values [1, 2], severity warn)]),\n" + " amount (type DOUBLE),\n status (type order_status),\n ),\n);\n\n" + "SELECT CAST(id AS INTEGER) AS order_id, amount, 'PAID' AS status\n" + 'FROM __source("raw_orders")\n' + ), + "models/marts/order_totals.sql": ( + 'MODEL (description "Order totals", materialized table);\n\n' + 'SELECT o.order_id, @doubled("o.amount") AS doubled_amount, o.status\n' + 'FROM __ref("stg_orders") AS o WHERE o.amount > 7\n' + ), + "models/marts/order_flags.sql": ( + 'MODEL (description "Order flags");\n\n' + 'SELECT o.order_id, o.missing_flag AS flag FROM __ref("stg_orders") AS o\n' + ), + "models/marts/flag_counts.sql": ( + 'MODEL (description "Flag counts");\n\nSELECT f.flag FROM __ref("order_flags") AS f\n' + ), + "tests/unit/test_order_totals.sql": ( + "TEST ();\n\nWITH\n__ref__stg_orders AS (\n" + " SELECT 1 AS order_id, 10.0 AS amount, 'PAID' AS status\n),\n" + "__expected__order_totals AS (\n" + " SELECT 1 AS order_id, 20.0 AS doubled_amount, 'PAID' AS status\n)\n" + "SELECT 1\n" + ), +} + + +_STAGING: str = "models/staging/stg_orders.sql" +_POPULATED_FACTS: tuple[str, ...] = ( + "project.sources", + "project.seeds", + "project.functions", + "tests.all", + "audits.all", + "declarations.enums", + "declarations.constants", + f"columns.declared:{_STAGING}", + f"columns.inferred:{_STAGING}", + f"audits.for_model:{_STAGING}", + "tests.for_model:models/marts/order_totals.sql", + "graph.dependents:models/marts/order_flags.sql", ) @@ -25,137 +98,67 @@ "test_case", ( HostPayloadFactsTestCase( - description="sources, refs, contracts, macros, declarations, tests, and audits", - project_files={ - "sqlbuild_project.toml": 'name = "orders"\nadapter = "duckdb"\n', - "sources/raw.yml": ( - "sources:\n" - " - name: raw_orders\n description: Raw orders.\n" - ' expression: "(SELECT 1 AS id, CAST(10.5 AS DOUBLE) AS amount)"\n' - " columns:\n - name: id\n type: INTEGER\n" - " - name: raw_customers\n description: Raw customers.\n" - ' expression: "(SELECT 1 AS customer_id)"\n' - ), - "constants/policy.sql": "CONSTANT (name minimum_amount, value 7);\n", - "enums/order_status.sql": "ENUM (name order_status, members [PAID, OPEN]);\n", - "models/marts/_sqlbuild/_macros/amounts.py": ( - 'def doubled(amount: str) -> str:\n return f"({amount}) * 2"\n' - ), - "models/staging/stg_orders.sql": ( - "MODEL (description 'Staged orders.',\n contract enforced,\n" - " columns (\n" - " order_id (type INTEGER, nullable false, audits [not_null]),\n" - " amount (type DOUBLE),\n status (type order_status),\n ),\n);\n\n" - "SELECT CAST(id AS INTEGER) AS order_id, amount, 'PAID' AS status\n" - 'FROM __source("raw_orders")\n' - ), - "models/marts/order_totals.sql": ( - 'MODEL (description "Order totals", materialized table);\n\n' - 'SELECT o.order_id, @doubled("o.amount") AS doubled_amount, o.status\n' - 'FROM __ref("stg_orders") AS o WHERE o.amount > 7\n' - ), - "models/marts/order_flags.sql": ( - 'MODEL (description "Order flags");\n\n' - 'SELECT o.order_id, o.missing_flag AS flag FROM __ref("stg_orders") AS o\n' - ), - "models/marts/flag_counts.sql": ( - 'MODEL (description "Flag counts");\n\n' - 'SELECT f.flag FROM __ref("order_flags") AS f\n' - ), - "hooks/sql/analyze_orders.sql": ( - 'HOOK (description "Analyze orders");\n\nSELECT 1\n' - ), - "hooks/python/order_hooks.py": ( - "from sqlbuild.hooks import hook\n\n" - "@hook\n" - "def log_orders(ctx):\n" - " '''Log orders.'''\n" - " ctx.execute_sql('INSERT INTO audit_log SELECT 1')\n" - ), - "python/loaders/order_rows.py": ( - "from sqlbuild.loaders import loader\n\n" - "@loader(columns=[{'name': 'order_id', 'type': 'INTEGER'}])\n" - "def fetch_orders(ctx):\n" - " '''Fetch orders.'''\n" - " return [{'order_id': 1}]\n" - ), - "materializations/copy_table.py": ( - "from sqlbuild.executor.custom.models import (\n" - " MaterializationContext,\n MaterializationResult,\n)\n\n\n" - "def materialize(ctx: MaterializationContext) -> MaterializationResult:\n" - ' ctx.execute_sql(f"CREATE TABLE {ctx.destination} AS {ctx.sql}")\n' - " return MaterializationResult(relation=ctx.destination)\n" - ), - "tests/scenarios/paid_orders.sql": ( - "SCENARIO (description 'Paid orders.');\n\n" - "WITH\n__source__raw_orders AS (SELECT 1 AS id, 10.5 AS amount),\n" - "__expected__stg_orders AS (\n" - " SELECT 1 AS order_id, 10.5 AS amount, 'PAID' AS status\n)\n" - "SELECT 1\n" - ), - "tests/unit/test_order_totals.sql": ( - "TEST ();\n\nWITH\n__ref__stg_orders AS (\n" - " SELECT 1 AS order_id, 10.0 AS amount\n),\n" - "__expected__order_totals AS (\n" - " SELECT 1 AS order_id, 20.0 AS doubled_amount, 'PAID' AS status\n)\n" - "SELECT 1\n" - ), - }, - expected_dropped_project_fields=frozenset( - { - "binding_catalog", - "diagnostics", - "external_sql_reference_resolver", - "hook_functions", - "loaded_macros", - "loader_functions", - "materialization_files", - "models", - "native_project", - "scope_index", - "sql_expansions", - "sql_hook_files", - "sql_scenarios", - "unmatched_literal_sql_relations", - } - ), - expected_dropped_model_fields=frozenset( - { - "authored_query_sql", - "binding_diagnostics", - "enum_columns", - "fast_lineage_columns", - "macro_deps", - "output_column_locations", - "references", - "unchecked_output_columns", - } - ), + description="native compiled project", + project_files=_PROJECT_FILES, + payload_source=native_project, + expected_populated_facts=_POPULATED_FACTS, + ), + HostPayloadFactsTestCase( + description="compiled project and model records", + project_files=_PROJECT_FILES, + payload_source=record_project, + expected_populated_facts=_POPULATED_FACTS, ), ), ids=lambda case: case.description, ) -def test_given_compiled_project_when_trimming_host_payload_then_every_fact_is_unchanged( +def test_given_compiled_project_when_decoding_host_payload_then_every_fact_is_unchanged( tmp_path: Path, test_case: HostPayloadFactsTestCase ) -> None: project: CompiledProjectView = compile_files(root=tmp_path, files=test_case.project_files) - trimmed: CompiledProjectView = custom_host_payload(project) + config: RulesConfig = RulesConfig() - full_digests: dict[FactKey, str | None] = fact_digests(project=project, root=tmp_path) - trimmed_digests: dict[FactKey, str | None] = fact_digests(project=trimmed, root=tmp_path) - - assert trimmed_digests == full_digests - assert tuple(map(public_model, trimmed.models)) == tuple(map(public_model, project.models)) - assert ( - dropped_fields(before=compiled_project_record(project), after=trimmed) - == test_case.expected_dropped_project_fields + decoded, decoded_config = decode_host_payload( + custom_host_payload(test_case.payload_source(project), config) + ) + expected: dict[str, object] = every_fact_value( + fact_views(project=ReferenceProject(project), root=tmp_path) ) - assert ( - frozenset().union( - *( - dropped_fields(before=model, after=trimmed_model) - for model, trimmed_model in zip(project.models, trimmed.models, strict=True) - ) - ) - == test_case.expected_dropped_model_fields + actual: dict[str, object] = every_fact_value(fact_views(project=decoded, root=tmp_path)) + + assert decoded_config == config + assert actual == expected + assert repr(actual) == repr(expected) + assert all(expected[key] for key in test_case.expected_populated_facts) + assert decoded.effective_target_name == project.effective_target_name + + +@pytest.mark.parametrize( + "test_case", + ( + HostPayloadLazyReadTestCase( + description="contract reads leave SQL, columns, tests and audits encoded", + project_files=_PROJECT_FILES, + expected_enforced=(False, False, False, True), + expected_encoded_model_facts=frozenset( + {"authored_sql", "query_sql", "declared_columns", "inferred_columns"} + ), + expected_encoded_project_facts=frozenset({"sql_tests", "audits", "sources"}), + ), + ), + ids=lambda case: case.description, +) +def test_given_rule_reading_one_fact_per_model_when_decoding_then_other_facts_stay_encoded( + tmp_path: Path, test_case: HostPayloadLazyReadTestCase +) -> None: + project: CompiledProjectView = compile_files(root=tmp_path, files=test_case.project_files) + + decoded: RuleFactProject = decode_host_payload(custom_host_payload(project, RulesConfig()))[0] + views: RuleFactViews = fact_views(project=decoded, root=tmp_path) + enforced: tuple[bool, ...] = tuple(map(views.contracts.enforced, views.project.models)) + + assert enforced == test_case.expected_enforced + assert all( + test_case.expected_encoded_model_facts.isdisjoint(vars(model)) for model in decoded.models ) + assert test_case.expected_encoded_project_facts.isdisjoint(vars(decoded)) 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 7f2dd7196c..5ba1f73053 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 @@ -411,7 +411,7 @@ def test_given_waiting_host_when_parent_exits_or_run_is_abandoned_then_wait_stop with pytest.raises(test_case.expected_error): custom_host.await_inputs( - path=tmp_path / "project.pickle", + path=tmp_path / "project.payload", cancel_markers=(tmp_path / "project.abandoned",), parent_pid=4242, ) @@ -425,7 +425,7 @@ def test_given_waiting_host_when_parent_exits_or_run_is_abandoned_then_wait_stop AwaitPublishedInputsTestCase( description="payload published while the parent runs", reported_parent_pid=4242, - existing_files=("project.pickle",), + existing_files=("project.payload",), expected_max_seconds=1.0, ), ), @@ -439,7 +439,7 @@ def test_given_published_payload_when_parent_still_runs_then_wait_returns( started: float = time.monotonic() custom_host.await_inputs( - path=tmp_path / "project.pickle", + path=tmp_path / "project.payload", cancel_markers=(tmp_path / "project.abandoned",), parent_pid=4242, )