From 8e46a4858e6f87e1a2cfdca91490717f4e1d269f Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 18 Oct 2026 02:41:51 +0100 Subject: [PATCH 1/3] perf: write the custom-rule host payload natively and decode facts lazily Custom-rule hosts received the whole compiled project as Python records pickled on the compiling thread, then unpickled every value before running any rule. The payload is now written by native code: model names, paths, keys, SQL text, config values, inferred columns and natively parsed header columns are stored without building Python objects, and only values that Python alone holds are pickled one by one. The host reads the header and decodes each fact the first time a rule reads it, so a rule touching one fact per model decodes only that fact. Fact views read the host project through narrow protocols; collections and declarations are read on first access. Read recording, cache keys and the host process lifecycle are unchanged. --- .../_helpers/model_config/header_metadata.rs | 75 +++- .../_helpers/rules/custom_evaluation.rs | 68 +--- .../_helpers/rules/custom_host_columns.rs | 254 ++++++++++++++ .../_helpers/rules/custom_host_payload.rs | 328 ++++++++++++++++++ .../bindings/_helpers/rules/custom_hosts.rs | 91 ++--- .../bindings/_helpers/rules/fact_replay.rs | 4 +- .../src/bindings/_helpers/rules/mod.rs | 2 + fensu.toml | 2 +- src/sqlbuild/_native.pyi | 4 +- .../_helpers/host/custom_evaluation.py | 15 +- .../rule_engine/_helpers/host/custom_host.py | 38 +- .../rule_engine/_helpers/host/host_payload.py | 189 ++++++++++ .../rule_engine/classes/audit_facts.py | 10 +- .../rule_engine/classes/column_facts.py | 11 +- .../rule_engine/classes/contract_facts.py | 10 +- .../rule_engine/classes/declaration_facts.py | 34 +- .../rule_engine/classes/graph_facts.py | 12 +- .../rule_engine/classes/model_test_facts.py | 11 +- .../rule_engine/classes/project_facts.py | 31 +- .../rule_engine/classes/project_tree.py | 30 +- .../rule_engine/classes/rule_context.py | 8 +- src/sqlbuild/rule_engine/classes/sql_facts.py | 12 +- src/sqlbuild/rule_engine/constants.py | 1 - src/sqlbuild/rule_engine/types.py | 85 ++++- .../sqlbuild/cli/commands/main/test_rules.py | 2 +- .../rule_engine/host_payload/_test_types.py | 20 +- .../rule_engine/host_payload/helpers.py | 178 +++++----- .../host_payload/test_host_payload_facts.py | 269 +++++++------- .../_helpers/host/test_custom_host_pool.py | 6 +- 29 files changed, 1321 insertions(+), 479 deletions(-) create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_columns.rs create mode 100644 crates/sqlbuild-python/src/bindings/_helpers/rules/custom_host_payload.rs create mode 100644 src/sqlbuild/rule_engine/_helpers/host/host_payload.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_rules.py b/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py index 2ea7a60274..fce7046cc0 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py @@ -1294,7 +1294,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( 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, ) From 6998c6f6b00ce6f549124d1245c8a4ae88ae65ec Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 18 Oct 2026 02:42:02 +0100 Subject: [PATCH 2/3] test: cover every rule fact, cache reuse and host outcomes through the native payload Rules reading every public fact report the compiled values cold and warm and rerun only the edited model's invocation; failing, host-exiting and slow rules keep their results, messages and payload cleanup. --- .../sqlbuild/cli/commands/main/_test_types.py | 21 ++ .../sqlbuild/cli/commands/main/test_rules.py | 228 ++++++++++++++++++ 2 files changed, 249 insertions(+) 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 fce7046cc0..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} @@ -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()) From a7b232c4a1a5812a07bedcb7ac29a2776e5b4e5d Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sun, 18 Oct 2026 02:48:55 +0100 Subject: [PATCH 3/3] docs: record performance measurement lessons from the cold compile work --- AGENTS.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) 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`