diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/errors.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/errors.rs new file mode 100644 index 0000000000..c2b2ea399e --- /dev/null +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/errors.rs @@ -0,0 +1,23 @@ +//! Error types owned by SQL-test planning helpers. + +use std::fmt::Display; + +use serde::ser; + +/// A parsed statement that `serde_json::to_value` could not represent. +#[derive(Debug)] +pub(crate) struct RelationMarkerError; + +impl Display for RelationMarkerError { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("value cannot be represented as JSON") + } +} + +impl std::error::Error for RelationMarkerError {} + +impl ser::Error for RelationMarkerError { + fn custom(_message: M) -> Self { + RelationMarkerError + } +} diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/mod.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/mod.rs index edada351e0..2efa1df628 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/mod.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/mod.rs @@ -1,9 +1,11 @@ pub(crate) mod cte_rename; pub(crate) mod cte_slices; pub(crate) mod cte_sql; +pub(crate) mod errors; pub(crate) mod expected_columns; pub(crate) mod extraction; pub(crate) mod helper_scope; pub(crate) mod markers; pub(crate) mod planning; +pub(crate) mod relation_markers; pub(crate) mod rendering; diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/planning.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/planning.rs index 8e477285cb..705ee2c626 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/planning.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/planning.rs @@ -8,12 +8,12 @@ use polyglot_sql::{Dialect, Expression}; use rayon::iter::{IntoParallelIterator, ParallelIterator}; use regex::Regex; use serde::{Deserialize, Serialize}; -use serde_json::Value; use crate::compiler::_helpers::sql_tests::cte_slices::SliceDialect; use crate::compiler::_helpers::sql_tests::cte_sql::{ cte_definition_sql, leading_with_prefix_end, with_leading_ctes, with_unique_ctes, }; +use crate::compiler::_helpers::sql_tests::errors::RelationMarkerError; use crate::compiler::_helpers::sql_tests::expected_columns::expected_columns; use crate::compiler::_helpers::sql_tests::helper_scope::{ ScopeGraph, helper_scope_ctes, merged_scoped_ctes, @@ -22,6 +22,7 @@ use crate::compiler::_helpers::sql_tests::markers::{ ProtectedRanges, marker_names, marker_names_in, replace_callable_markers, replace_dbt_ref_markers, replace_named_markers, }; +use crate::compiler::_helpers::sql_tests::relation_markers::relation_marker_calls; use crate::compiler::_helpers::sql_tests::rendering::{ AssertionStep, ChainStep, RenderRequest, render_comparison_sql, render_dialect, rendered_chain_steps, @@ -49,7 +50,7 @@ const ASSERT_PREFIX: &str = "__assert__"; const REF_FUNCTION: &str = "__ref"; const SOURCE_FUNCTION: &str = "__source"; const SEED_FUNCTION: &str = "__seed"; -const DBT_REF_FUNCTION: &str = "__dbt_ref"; +pub(crate) const DBT_REF_FUNCTION: &str = "__dbt_ref"; const UDF_FUNCTION: &str = "__udf"; const TABLE_FUNCTION: &str = "__table_fn"; @@ -1304,13 +1305,13 @@ fn analysis_template( return None; } let expression = statements.remove(0); - let value = match serde_json::to_value(&expression) { - Ok(value) => value, - Err(_) => return None, + let marker_calls = match relation_marker_calls(&expression) { + Ok(marker_calls) => marker_calls, + Err(RelationMarkerError) => return None, }; Some(AnalysisTemplate { existing_cte_names: top_level_cte_names(&expression), - marker_calls: relation_marker_calls(&value), + marker_calls, }) }) }) @@ -1484,73 +1485,6 @@ fn resolve_function_calls( Ok(result) } -fn relation_marker_calls(value: &Value) -> Vec<(String, String)> { - fn marker(expression: &Value) -> Option<(String, String)> { - let object = expression.as_object()?; - if let Some(alias) = object.get("alias") { - return marker(alias.get("this")?); - } - let function = object.get("function")?.as_object()?; - let function_name = function.get("name")?.as_str()?.to_ascii_lowercase(); - let args = function.get("args")?.as_array()?; - let arg_name = |arg: &Value| { - arg.get("column")? - .get("name")? - .get("name")? - .as_str() - .map(str::to_string) - }; - let referenced_name = if function_name == DBT_REF_FUNCTION { - match args.as_slice() { - [first] => arg_name(first)?, - [first, second] => format!("{}__{}", arg_name(first)?, arg_name(second)?), - _ => return None, - } - } else { - let [first] = args.as_slice() else { - return None; - }; - arg_name(first)? - }; - Some((function_name, referenced_name)) - } - - fn walk(value: &Value, calls: &mut Vec<(String, String)>) { - match value { - Value::Object(object) => { - if let Some(expressions) = object - .get("from") - .and_then(|from| from.get("expressions")) - .and_then(Value::as_array) - { - calls.extend(expressions.iter().filter_map(marker)); - } - if let Some(joins) = object.get("joins").and_then(Value::as_array) { - calls.extend( - joins - .iter() - .filter_map(|join| join.get("this")) - .filter_map(marker), - ); - } - for child in object.values() { - walk(child, calls); - } - } - Value::Array(values) => { - for child in values { - walk(child, calls); - } - } - _ => {} - } - } - - let mut calls: Vec<(String, String)> = Vec::new(); - walk(value, &mut calls); - calls -} - fn top_level_cte_names(expression: &Expression) -> HashSet { let Expression::Select(select) = expression else { return HashSet::new(); diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/relation_markers.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/relation_markers.rs new file mode 100644 index 0000000000..c7b3dea5d6 --- /dev/null +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/relation_markers.rs @@ -0,0 +1,809 @@ +//! Relation marker calls read from a parsed statement's serialized shape without a JSON tree. + +use std::cell::RefCell; + +use serde::Serialize; +use serde::ser::{ + self, SerializeMap, SerializeSeq, SerializeStruct, SerializeStructVariant, SerializeTuple, + SerializeTupleStruct, SerializeTupleVariant, +}; + +use crate::compiler::_helpers::sql_tests::errors::RelationMarkerError; +use crate::compiler::_helpers::sql_tests::planning::DBT_REF_FUNCTION; + +pub(crate) type MarkerCall = (String, String); + +/// Return the marker relations a walk of `serde_json::to_value(value)` finds, in the same order. +pub(crate) fn relation_marker_calls(value: &T) -> Result> { + let stack: RefCell> = RefCell::new(Vec::new()); + let out = value.serialize(Collector { + need: Need::Walk, + stack: &stack, + })?; + Ok(out.calls) +} + +/// What the parent reads from a value besides the marker calls found beneath it. +#[derive(Clone, Copy, PartialEq, Eq)] +enum Need { + Walk, + From, + FromExpressions, + Joins, + Marker, + ThisMarker, + Function, + Args, + Arg, + Column, + Name, + Str, +} + +impl Need { + fn field(self, key: &str) -> Need { + match (self, key) { + (_, "from") => Need::From, + (_, "joins") => Need::Joins, + (Need::From, "expressions") => Need::FromExpressions, + (Need::Marker, "alias") => Need::ThisMarker, + (Need::Marker, "function") => Need::Function, + (Need::ThisMarker, "this") => Need::Marker, + (Need::Function, "name") | (Need::Name, "name") => Need::Str, + (Need::Function, "args") => Need::Args, + (Need::Arg, "column") => Need::Column, + (Need::Column, "name") => Need::Name, + _ => Need::Walk, + } + } + + fn element(self) -> Need { + match self { + Need::FromExpressions => Need::Marker, + Need::Joins => Need::ThisMarker, + Need::Args => Need::Arg, + _ => Need::Walk, + } + } +} + +#[derive(Default)] +enum Info { + #[default] + None, + Str(String), + Marker(MarkerCall), + Markers(Vec>), + Names(Vec>), + Function(Option, Option>>), +} + +impl Info { + fn into_str(self) -> Option { + match self { + Info::Str(value) => Some(value), + _ => None, + } + } + + fn into_marker(self) -> Option { + match self { + Info::Marker(value) => Some(value), + _ => None, + } + } +} + +#[derive(Default)] +struct Out { + calls: Vec, + info: Info, +} + +impl Out { + fn leaf() -> Self { + Self::default() + } +} + +type Result = std::result::Result; + +enum Key { + Static(&'static str), + Owned(String), +} + +impl Key { + fn as_str(&self) -> &str { + match self { + Key::Static(value) => value, + Key::Owned(value) => value, + } + } +} + +type Stack = RefCell>; + +/// Pop the entries pushed since `start` as one object; a repeated key keeps its first position. +fn finish_object(need: Need, stack: &Stack, start: usize, may_repeat_keys: bool) -> Out { + let mut stack = stack.borrow_mut(); + let mut index = start + 1; + while may_repeat_keys && index < stack.len() { + let name = stack[index].0.as_str(); + match stack[start..index] + .iter() + .position(|(key, _)| key.as_str() == name) + { + Some(first) => { + let (_, out) = stack.remove(index); + stack[start + first].1 = out; + } + None => index += 1, + } + } + let entries = &mut stack[start..]; + let mut take = |name: &str| -> Option { + entries + .iter_mut() + .find(|(key, _)| key.as_str() == name) + .map(|(_, out)| std::mem::take(&mut out.info)) + }; + let mut calls: Vec = Vec::new(); + for name in ["from", "joins"] { + if let Some(Info::Markers(markers)) = take(name) { + calls.extend(markers.into_iter().flatten()); + } + } + let info = match need { + Need::Marker => match take("alias") { + Some(alias) => alias.into_marker().map_or(Info::None, Info::Marker), + None => match take("function") { + Some(Info::Function(Some(name), Some(args))) => { + function_marker(&name, args).map_or(Info::None, Info::Marker) + } + _ => Info::None, + }, + }, + Need::ThisMarker => take("this").unwrap_or_default(), + Need::From => match take("expressions") { + Some(markers @ Info::Markers(_)) => markers, + _ => Info::None, + }, + Need::Function => Info::Function( + take("name").and_then(Info::into_str), + match take("args") { + Some(Info::Names(names)) => Some(names), + _ => None, + }, + ), + Need::Arg => take("column").unwrap_or_default(), + Need::Column | Need::Name => take("name").unwrap_or_default(), + Need::Walk | Need::FromExpressions | Need::Joins | Need::Args | Need::Str => Info::None, + }; + for (_, out) in stack.drain(start..) { + calls.extend(out.calls); + } + Out { calls, info } +} + +/// One bit per key shape, so repeated keys are checked exactly only when a bit repeats. +fn key_bit(key: &str) -> u64 { + let bytes = key.as_bytes(); + let shape = bytes.len() + ^ usize::from(bytes.first().copied().unwrap_or(0)).rotate_left(3) + ^ usize::from(bytes.last().copied().unwrap_or(0)).rotate_left(7); + 1u64 << (shape % 64) +} + +fn function_marker(name: &str, args: Vec>) -> Option { + let function_name = name.to_ascii_lowercase(); + let referenced_name = if function_name == DBT_REF_FUNCTION { + match args.as_slice() { + [Some(first)] => first.clone(), + [Some(first), Some(second)] => format!("{first}__{second}"), + _ => return None, + } + } else { + let [Some(first)] = args.as_slice() else { + return None; + }; + first.clone() + }; + Some((function_name, referenced_name)) +} + +fn finish_array(need: Need, stack: &Stack, start: usize) -> Out { + let mut calls: Vec = Vec::new(); + let mut stack = stack.borrow_mut(); + let info = match need { + Need::FromExpressions | Need::Joins => Info::Markers( + stack[start..] + .iter_mut() + .map(|(_, out)| std::mem::take(&mut out.info).into_marker()) + .collect(), + ), + Need::Args => Info::Names( + stack[start..] + .iter_mut() + .map(|(_, out)| std::mem::take(&mut out.info).into_str()) + .collect(), + ), + _ => Info::None, + }; + for (_, out) in stack.drain(start..) { + calls.extend(out.calls); + } + Out { calls, info } +} + +#[derive(Clone, Copy)] +struct Collector<'a> { + need: Need, + stack: &'a Stack, +} + +impl<'a> Collector<'a> { + fn string(self, value: impl FnOnce() -> String) -> Out { + Out { + calls: Vec::new(), + info: if self.need == Need::Str { + Info::Str(value()) + } else { + Info::None + }, + } + } + + fn array(self, need: Need) -> ArrayFrame<'a> { + ArrayFrame { + need, + stack: self.stack, + start: self.stack.borrow().len(), + } + } + + /// Open `{variant: payload}`; the outer object is empty until the payload closes. + fn variant( + self, + variant: &'static str, + payload: fn(Self, Need) -> F, + ) -> VariantFrame<'a, F> { + VariantFrame { + outer: self.object(self.need), + variant, + inner: payload(self, self.need.field(variant)), + } + } + + fn object(self, need: Need) -> ObjectFrame<'a> { + ObjectFrame { + need, + stack: self.stack, + start: self.stack.borrow().len(), + key_bits: 0, + may_repeat_keys: false, + pending_key: None, + } + } +} + +struct ArrayFrame<'a> { + need: Need, + stack: &'a Stack, + start: usize, +} + +impl ArrayFrame<'_> { + fn push(&mut self, value: &T) -> Result<()> { + let out = value.serialize(Collector { + need: self.need.element(), + stack: self.stack, + })?; + self.stack.borrow_mut().push((Key::Static(""), out)); + Ok(()) + } + + fn finish(self) -> Out { + finish_array(self.need, self.stack, self.start) + } +} + +struct ObjectFrame<'a> { + need: Need, + stack: &'a Stack, + start: usize, + key_bits: u64, + may_repeat_keys: bool, + pending_key: Option, +} + +impl ObjectFrame<'_> { + fn push(&mut self, key: Key, value: &T) -> Result<()> { + let out = value.serialize(Collector { + need: self.need.field(key.as_str()), + stack: self.stack, + })?; + let bit = key_bit(key.as_str()); + self.may_repeat_keys |= self.key_bits & bit != 0; + self.key_bits |= bit; + self.stack.borrow_mut().push((key, out)); + Ok(()) + } + + fn finish(self) -> Out { + finish_object(self.need, self.stack, self.start, self.may_repeat_keys) + } +} + +/// A variant wrapper object `{variant: inner}` around an array or object payload. +struct VariantFrame<'a, F> { + outer: ObjectFrame<'a>, + variant: &'static str, + inner: F, +} + +impl VariantFrame<'_, F> { + fn finish(self, inner: Out) -> Out { + let outer = self.outer; + outer + .stack + .borrow_mut() + .push((Key::Static(self.variant), inner)); + outer.finish() + } +} + +impl<'a> ser::Serializer for Collector<'a> { + type Ok = Out; + type Error = RelationMarkerError; + type SerializeSeq = ArrayFrame<'a>; + type SerializeTuple = ArrayFrame<'a>; + type SerializeTupleStruct = ArrayFrame<'a>; + type SerializeTupleVariant = VariantFrame<'a, ArrayFrame<'a>>; + type SerializeMap = ObjectFrame<'a>; + type SerializeStruct = ObjectFrame<'a>; + type SerializeStructVariant = VariantFrame<'a, ObjectFrame<'a>>; + + fn serialize_bool(self, _value: bool) -> Result { + Ok(Out::leaf()) + } + + fn serialize_i8(self, _value: i8) -> Result { + Ok(Out::leaf()) + } + + fn serialize_i16(self, _value: i16) -> Result { + Ok(Out::leaf()) + } + + fn serialize_i32(self, _value: i32) -> Result { + Ok(Out::leaf()) + } + + fn serialize_i64(self, _value: i64) -> Result { + Ok(Out::leaf()) + } + + fn serialize_i128(self, value: i128) -> Result { + if (i128::from(i64::MIN)..=i128::from(u64::MAX)).contains(&value) { + Ok(Out::leaf()) + } else { + Err(RelationMarkerError) + } + } + + fn serialize_u8(self, _value: u8) -> Result { + Ok(Out::leaf()) + } + + fn serialize_u16(self, _value: u16) -> Result { + Ok(Out::leaf()) + } + + fn serialize_u32(self, _value: u32) -> Result { + Ok(Out::leaf()) + } + + fn serialize_u64(self, _value: u64) -> Result { + Ok(Out::leaf()) + } + + fn serialize_u128(self, value: u128) -> Result { + if value <= u128::from(u64::MAX) { + Ok(Out::leaf()) + } else { + Err(RelationMarkerError) + } + } + + fn serialize_f32(self, _value: f32) -> Result { + Ok(Out::leaf()) + } + + fn serialize_f64(self, _value: f64) -> Result { + Ok(Out::leaf()) + } + + fn serialize_char(self, value: char) -> Result { + Ok(self.string(|| value.to_string())) + } + + fn serialize_str(self, value: &str) -> Result { + Ok(self.string(|| value.to_string())) + } + + fn serialize_bytes(self, _value: &[u8]) -> Result { + Ok(Out::leaf()) + } + + fn serialize_none(self) -> Result { + Ok(Out::leaf()) + } + + fn serialize_some(self, value: &T) -> Result { + value.serialize(self) + } + + fn serialize_unit(self) -> Result { + Ok(Out::leaf()) + } + + fn serialize_unit_struct(self, _name: &'static str) -> Result { + Ok(Out::leaf()) + } + + fn serialize_unit_variant( + self, + _name: &'static str, + _variant_index: u32, + variant: &'static str, + ) -> Result { + Ok(self.string(|| variant.to_string())) + } + + fn serialize_newtype_struct( + self, + _name: &'static str, + value: &T, + ) -> Result { + value.serialize(self) + } + + fn serialize_newtype_variant( + self, + _name: &'static str, + _variant_index: u32, + variant: &'static str, + value: &T, + ) -> Result { + let mut frame = self.object(self.need); + frame.push(Key::Static(variant), value)?; + Ok(frame.finish()) + } + + fn serialize_seq(self, _len: Option) -> Result> { + Ok(self.array(self.need)) + } + + fn serialize_tuple(self, _len: usize) -> Result> { + Ok(self.array(self.need)) + } + + fn serialize_tuple_struct(self, _name: &'static str, _len: usize) -> Result> { + Ok(self.array(self.need)) + } + + fn serialize_tuple_variant( + self, + _name: &'static str, + _variant_index: u32, + variant: &'static str, + _len: usize, + ) -> Result>> { + Ok(self.variant(variant, Collector::array)) + } + + fn serialize_map(self, _len: Option) -> Result> { + Ok(self.object(self.need)) + } + + fn serialize_struct(self, _name: &'static str, _len: usize) -> Result> { + Ok(self.object(self.need)) + } + + fn serialize_struct_variant( + self, + _name: &'static str, + _variant_index: u32, + variant: &'static str, + _len: usize, + ) -> Result>> { + Ok(self.variant(variant, Collector::object)) + } +} + +impl SerializeSeq for ArrayFrame<'_> { + type Ok = Out; + type Error = RelationMarkerError; + + fn serialize_element(&mut self, value: &T) -> Result<()> { + self.push(value) + } + + fn end(self) -> Result { + Ok(self.finish()) + } +} + +impl SerializeTuple for ArrayFrame<'_> { + type Ok = Out; + type Error = RelationMarkerError; + + fn serialize_element(&mut self, value: &T) -> Result<()> { + self.push(value) + } + + fn end(self) -> Result { + Ok(self.finish()) + } +} + +impl SerializeTupleStruct for ArrayFrame<'_> { + type Ok = Out; + type Error = RelationMarkerError; + + fn serialize_field(&mut self, value: &T) -> Result<()> { + self.push(value) + } + + fn end(self) -> Result { + Ok(self.finish()) + } +} + +impl SerializeTupleVariant for VariantFrame<'_, ArrayFrame<'_>> { + type Ok = Out; + type Error = RelationMarkerError; + + fn serialize_field(&mut self, value: &T) -> Result<()> { + self.inner.push(value) + } + + fn end(self) -> Result { + let inner = finish_array(self.inner.need, self.inner.stack, self.inner.start); + Ok(self.finish(inner)) + } +} + +impl SerializeMap for ObjectFrame<'_> { + type Ok = Out; + type Error = RelationMarkerError; + + fn serialize_key(&mut self, key: &T) -> Result<()> { + self.pending_key = Some(Key::Owned(key.serialize(MapKeyCollector)?)); + Ok(()) + } + + fn serialize_value(&mut self, value: &T) -> Result<()> { + let key = self.pending_key.take().ok_or(RelationMarkerError)?; + self.push(key, value) + } + + fn end(self) -> Result { + Ok(self.finish()) + } +} + +impl SerializeStruct for ObjectFrame<'_> { + type Ok = Out; + type Error = RelationMarkerError; + + fn serialize_field( + &mut self, + key: &'static str, + value: &T, + ) -> Result<()> { + self.push(Key::Static(key), value) + } + + fn end(self) -> Result { + Ok(self.finish()) + } +} + +impl SerializeStructVariant for VariantFrame<'_, ObjectFrame<'_>> { + type Ok = Out; + type Error = RelationMarkerError; + + fn serialize_field( + &mut self, + key: &'static str, + value: &T, + ) -> Result<()> { + self.inner.push(Key::Static(key), value) + } + + fn end(self) -> Result { + let inner = finish_object( + self.inner.need, + self.inner.stack, + self.inner.start, + self.inner.may_repeat_keys, + ); + Ok(self.finish(inner)) + } +} + +/// Map keys as `serde_json` stringifies them; unsupported key types fail as they do there. +struct MapKeyCollector; + +impl ser::Serializer for MapKeyCollector { + type Ok = String; + type Error = RelationMarkerError; + type SerializeSeq = ser::Impossible; + type SerializeTuple = ser::Impossible; + type SerializeTupleStruct = ser::Impossible; + type SerializeTupleVariant = ser::Impossible; + type SerializeMap = ser::Impossible; + type SerializeStruct = ser::Impossible; + type SerializeStructVariant = ser::Impossible; + + fn serialize_bool(self, value: bool) -> Result { + Ok(value.to_string()) + } + + fn serialize_i8(self, value: i8) -> Result { + Ok(value.to_string()) + } + + fn serialize_i16(self, value: i16) -> Result { + Ok(value.to_string()) + } + + fn serialize_i32(self, value: i32) -> Result { + Ok(value.to_string()) + } + + fn serialize_i64(self, value: i64) -> Result { + Ok(value.to_string()) + } + + fn serialize_i128(self, value: i128) -> Result { + Ok(value.to_string()) + } + + fn serialize_u8(self, value: u8) -> Result { + Ok(value.to_string()) + } + + fn serialize_u16(self, value: u16) -> Result { + Ok(value.to_string()) + } + + fn serialize_u32(self, value: u32) -> Result { + Ok(value.to_string()) + } + + fn serialize_u64(self, value: u64) -> Result { + Ok(value.to_string()) + } + + fn serialize_u128(self, value: u128) -> Result { + Ok(value.to_string()) + } + + fn serialize_f32(self, value: f32) -> Result { + if value.is_finite() { + Ok(format!("{value:?}")) + } else { + Err(RelationMarkerError) + } + } + + fn serialize_f64(self, value: f64) -> Result { + if value.is_finite() { + Ok(format!("{value:?}")) + } else { + Err(RelationMarkerError) + } + } + + fn serialize_char(self, value: char) -> Result { + Ok(value.to_string()) + } + + fn serialize_str(self, value: &str) -> Result { + Ok(value.to_string()) + } + + fn serialize_bytes(self, _value: &[u8]) -> Result { + Err(RelationMarkerError) + } + + fn serialize_none(self) -> Result { + Err(RelationMarkerError) + } + + fn serialize_some(self, _value: &T) -> Result { + Err(RelationMarkerError) + } + + fn serialize_unit(self) -> Result { + Err(RelationMarkerError) + } + + fn serialize_unit_struct(self, _name: &'static str) -> Result { + Err(RelationMarkerError) + } + + fn serialize_unit_variant( + self, + _name: &'static str, + _variant_index: u32, + variant: &'static str, + ) -> Result { + Ok(variant.to_string()) + } + + fn serialize_newtype_struct( + self, + _name: &'static str, + value: &T, + ) -> Result { + value.serialize(self) + } + + fn serialize_newtype_variant( + self, + _name: &'static str, + _variant_index: u32, + _variant: &'static str, + _value: &T, + ) -> Result { + Err(RelationMarkerError) + } + + fn serialize_seq(self, _len: Option) -> Result { + Err(RelationMarkerError) + } + + fn serialize_tuple(self, _len: usize) -> Result { + Err(RelationMarkerError) + } + + fn serialize_tuple_struct( + self, + _name: &'static str, + _len: usize, + ) -> Result { + Err(RelationMarkerError) + } + + fn serialize_tuple_variant( + self, + _name: &'static str, + _variant_index: u32, + _variant: &'static str, + _len: usize, + ) -> Result { + Err(RelationMarkerError) + } + + fn serialize_map(self, _len: Option) -> Result { + Err(RelationMarkerError) + } + + fn serialize_struct(self, _name: &'static str, _len: usize) -> Result { + Err(RelationMarkerError) + } + + fn serialize_struct_variant( + self, + _name: &'static str, + _variant_index: u32, + _variant: &'static str, + _len: usize, + ) -> Result { + Err(RelationMarkerError) + } +} diff --git a/crates/sqlbuild-rules-native/src/compiler/tests.rs b/crates/sqlbuild-rules-native/src/compiler/tests.rs index 4888371b81..1a3a7a928d 100644 --- a/crates/sqlbuild-rules-native/src/compiler/tests.rs +++ b/crates/sqlbuild-rules-native/src/compiler/tests.rs @@ -18,6 +18,8 @@ mod sql_test_extraction; mod sql_test_markers; #[path = "tests/test_sql_test_planning.rs"] mod sql_test_planning; +#[path = "tests/test_sql_test_relation_markers.rs"] +mod sql_test_relation_markers; #[path = "tests/test_sql_test_rendering.rs"] mod sql_test_rendering; #[path = "tests/test_types.rs"] diff --git a/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs b/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs index fcfaeb3e5a..2b6f36eb73 100644 --- a/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs +++ b/crates/sqlbuild-rules-native/src/compiler/tests/helpers.rs @@ -2,6 +2,7 @@ use std::collections::HashMap; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Barrier, Mutex}; +use polyglot_sql::Dialect; use serde_json::{Value, json}; use crate::compiler::_helpers::model_headers::tokenization::{ @@ -11,6 +12,7 @@ use crate::compiler::_helpers::sql_interpolation::substitution::{ FALLBACK, SUBSTITUTED, UNCHANGED, substitute_batch, }; use crate::compiler::_helpers::sql_references::extraction::extract; +use crate::compiler::_helpers::sql_tests::relation_markers::relation_marker_calls; use crate::compiler::main::sql_test_extraction::extract_batch_json; use crate::compiler::models::AuthoredValue; @@ -1657,3 +1659,133 @@ fn generic_lexical_syntax() -> Value { "lineCommentPrefixes": ["--"] }) } + +type MarkerCalls = Option>; + +/// The previous JSON-tree walk, kept as the oracle the streaming collector must match. +fn json_relation_marker_calls(value: &Value) -> Vec<(String, String)> { + let from_markers = value + .get("from") + .and_then(|from| from.get("expressions")) + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(json_marker); + let join_markers = value + .get("joins") + .and_then(Value::as_array) + .into_iter() + .flatten() + .filter_map(|join| join.get("this")) + .filter_map(json_marker); + let children: Vec<&Value> = value + .as_object() + .map(|object| object.values().collect()) + .or_else(|| value.as_array().map(|values| values.iter().collect())) + .unwrap_or_default(); + from_markers + .chain(join_markers) + .chain(children.into_iter().flat_map(json_relation_marker_calls)) + .collect() +} + +fn json_marker(expression: &Value) -> Option<(String, String)> { + let object = expression.as_object()?; + object.get("alias").map_or_else( + || json_function_marker(object), + |alias| alias.get("this").and_then(json_marker), + ) +} + +fn json_function_marker(object: &serde_json::Map) -> Option<(String, String)> { + let function = object.get("function")?.as_object()?; + let function_name = function.get("name")?.as_str()?.to_ascii_lowercase(); + let names: Vec = function + .get("args")? + .as_array()? + .iter() + .map(|arg| { + arg.get("column")? + .get("name")? + .get("name")? + .as_str() + .map(str::to_string) + }) + .collect::>()?; + let referenced_name = (function_name == "__dbt_ref" && names.len() == 2) + .then(|| names.join("__")) + .or_else(|| (names.len() == 1).then(|| names[0].clone()))?; + Some((function_name, referenced_name)) +} + +/// Owned marker calls for comparison with collected calls. +pub(crate) fn owned_marker_calls(calls: &[(&str, &str)]) -> MarkerCalls { + Some( + calls + .iter() + .map(|(function, name)| ((*function).to_string(), (*name).to_string())) + .collect(), + ) +} + +/// Parse one statement and return the streaming collector's and the JSON oracle's marker calls. +pub(crate) fn relation_markers_and_json_oracle( + dialect: &str, + sql: &str, +) -> (MarkerCalls, MarkerCalls) { + let mut statements = Dialect::get_by_name(dialect) + .expect("test dialect must exist") + .parse(sql) + .expect("test SQL must parse"); + assert_eq!(statements.len(), 1, "{sql}"); + let expression = statements.remove(0); + ( + relation_marker_calls(&expression).ok(), + serde_json::to_value(&expression) + .ok() + .map(|value| json_relation_marker_calls(&value)), + ) +} + +/// Return every generated relation shape whose collected markers differ from the JSON oracle. +pub(crate) fn relation_marker_oracle_mismatches(dialect: &str) -> Vec { + let relations: &[&str] = &[ + "__ref(\"orders\")", + "__source(\"raw_orders\") AS s", + "__dbt_ref(\"shop\", \"customers\") c", + "(SELECT * FROM __seed(\"regions\") JOIN __ref(\"stores\") ON TRUE) AS d", + "LATERAL (SELECT * FROM __ref(\"items\")) AS l", + "__udf(\"fn_orders\")", + "plain_table", + ]; + let wrappers = [ + "SELECT * FROM {a} JOIN {b} ON TRUE", + "WITH x AS (SELECT * FROM {a}) SELECT (SELECT 1 FROM {b}) AS v FROM x", + "SELECT * FROM {a} WHERE EXISTS (SELECT 1 FROM {b})", + "SELECT * FROM {a} UNION SELECT * FROM {b}", + "SELECT * FROM {a}, {b}", + ]; + let parser = Dialect::get_by_name(dialect).expect("test dialect must exist"); + wrappers + .iter() + .flat_map(|wrapper| { + relations.iter().flat_map(move |first| { + relations + .iter() + .map(move |second| wrapper.replace("{a}", first).replace("{b}", second)) + }) + }) + .filter(|sql| { + parser + .parse(sql) + .unwrap_or_default() + .iter() + .any(|expression| { + relation_marker_calls(expression).ok() + != serde_json::to_value(expression) + .ok() + .map(|value| json_relation_marker_calls(&value)) + }) + }) + .collect() +} diff --git a/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_relation_markers.rs b/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_relation_markers.rs new file mode 100644 index 0000000000..d1f123948a --- /dev/null +++ b/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_relation_markers.rs @@ -0,0 +1,111 @@ +use crate::compiler::tests::helpers::{ + owned_marker_calls, relation_marker_oracle_mismatches, relation_markers_and_json_oracle, +}; +use crate::compiler::tests::test_types::{RelationMarkerOracleTestCase, RelationMarkersTestCase}; + +#[test] +fn given_parsed_statements_when_collecting_relation_markers_then_json_walk_order_is_reproduced() { + let test_cases = [ + RelationMarkersTestCase { + description: "plain reference", + dialect: "duckdb", + sql: "SELECT * FROM __ref(\"orders\")", + expected_calls: &[("__ref", "orders")], + }, + RelationMarkersTestCase { + description: "aliased joins keep from before joins", + dialect: "snowflake", + sql: "SELECT o.id FROM __ref(\"orders\") AS o JOIN __source(\"raw_customers\") AS c ON o.id = c.id LEFT JOIN __seed(\"regions\") r ON c.region = r.id", + expected_calls: &[ + ("__ref", "orders"), + ("__source", "raw_customers"), + ("__seed", "regions"), + ], + }, + RelationMarkersTestCase { + description: "nested relations follow serialized field order", + dialect: "duckdb", + sql: "WITH recent AS (SELECT * FROM __ref(\"orders\")) SELECT (SELECT COUNT(*) FROM __ref(\"returns\")) AS returned FROM recent JOIN (SELECT * FROM __ref(\"customers\")) AS c ON TRUE WHERE EXISTS (SELECT 1 FROM __source(\"raw_events\"))", + expected_calls: &[ + ("__ref", "returns"), + ("__ref", "customers"), + ("__source", "raw_events"), + ("__ref", "orders"), + ], + }, + RelationMarkersTestCase { + description: "dbt references with one and two arguments", + dialect: "duckdb", + sql: "SELECT * FROM __dbt_ref(\"orders\") JOIN __dbt_ref(\"shop\", \"customers\") ON TRUE", + expected_calls: &[("__dbt_ref", "orders"), ("__dbt_ref", "shop__customers")], + }, + RelationMarkersTestCase { + description: "upper-case names are folded and unknown functions kept", + dialect: "duckdb", + sql: "SELECT * FROM __REF(\"Orders\") JOIN generate_rows(amount) ON TRUE", + expected_calls: &[("__ref", "Orders"), ("generate_rows", "amount")], + }, + RelationMarkersTestCase { + description: "calls with several arguments or outside relations are ignored", + dialect: "duckdb", + sql: "SELECT __ref(\"orders\") AS value FROM __table_fn(\"orders\", 2) WHERE id IN (SELECT id FROM t)", + expected_calls: &[], + }, + RelationMarkersTestCase { + description: "set operations walk each branch", + dialect: "postgres", + sql: "SELECT id FROM __ref(\"orders\") UNION ALL SELECT id FROM __ref(\"archived_orders\")", + expected_calls: &[("__ref", "orders"), ("__ref", "archived_orders")], + }, + ]; + for test_case in test_cases { + let (collected, oracle) = + relation_markers_and_json_oracle(test_case.dialect, test_case.sql); + assert_eq!( + collected, + owned_marker_calls(test_case.expected_calls), + "{}", + test_case.description + ); + assert_eq!( + oracle, + owned_marker_calls(test_case.expected_calls), + "{}", + test_case.description + ); + } +} + +#[test] +fn given_generated_relation_shapes_when_collecting_markers_then_json_oracle_matches() { + let test_cases = [ + RelationMarkerOracleTestCase { + description: "duckdb relation shapes", + dialect: "duckdb", + expected_mismatches: &[], + }, + RelationMarkerOracleTestCase { + description: "snowflake relation shapes", + dialect: "snowflake", + expected_mismatches: &[], + }, + RelationMarkerOracleTestCase { + description: "postgres relation shapes", + dialect: "postgres", + expected_mismatches: &[], + }, + RelationMarkerOracleTestCase { + description: "bigquery relation shapes", + dialect: "bigquery", + expected_mismatches: &[], + }, + ]; + for test_case in test_cases { + assert_eq!( + relation_marker_oracle_mismatches(test_case.dialect), + test_case.expected_mismatches, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-rules-native/src/compiler/tests/test_types.rs b/crates/sqlbuild-rules-native/src/compiler/tests/test_types.rs index 3f96531f7c..dff917362b 100644 --- a/crates/sqlbuild-rules-native/src/compiler/tests/test_types.rs +++ b/crates/sqlbuild-rules-native/src/compiler/tests/test_types.rs @@ -73,3 +73,16 @@ pub(crate) struct ProtectedRangeTestCase { pub(crate) ranges: &'static [(usize, usize)], pub(crate) expected_protected: &'static [usize], } + +pub(crate) struct RelationMarkersTestCase { + pub(crate) description: &'static str, + pub(crate) dialect: &'static str, + pub(crate) sql: &'static str, + pub(crate) expected_calls: &'static [(&'static str, &'static str)], +} + +pub(crate) struct RelationMarkerOracleTestCase { + pub(crate) description: &'static str, + pub(crate) dialect: &'static str, + pub(crate) expected_mismatches: &'static [&'static str], +} diff --git a/src/sqlbuild/cli/commands/_helpers/compile/pipeline.py b/src/sqlbuild/cli/commands/_helpers/compile/pipeline.py index 72b2548e2a..ae4817b5db 100644 --- a/src/sqlbuild/cli/commands/_helpers/compile/pipeline.py +++ b/src/sqlbuild/cli/commands/_helpers/compile/pipeline.py @@ -201,13 +201,12 @@ def _analyze_compile_project( ) if not graph_failed: rules_config: RulesConfig = load_rules_config(project_dir=project_dir) - if ( - prepared_artifacts is not None - and rules_config.select - and not any(item.is_error for item in core_diagnostics) - ): + if prepared_artifacts is not None and not any(item.is_error for item in core_diagnostics): prepared_artifacts.start( - project=graph.project, adapter=adapter, target_dir=project_dir / "target" + project=graph.project, + adapter=adapter, + target_dir=project_dir / "target", + stage_artifacts=bool(rules_config.select), ) _ = start_compile_phase(status=status, message="Evaluating built-in and custom rules...") rules_start: float = time.monotonic() diff --git a/src/sqlbuild/cli/commands/_helpers/compile/target_writer.py b/src/sqlbuild/cli/commands/_helpers/compile/target_writer.py index 0b07b618f5..aa95dc2d60 100644 --- a/src/sqlbuild/cli/commands/_helpers/compile/target_writer.py +++ b/src/sqlbuild/cli/commands/_helpers/compile/target_writer.py @@ -6,6 +6,7 @@ import json import os import uuid +from collections.abc import Callable from pathlib import Path from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter @@ -20,6 +21,8 @@ ) from sqlbuild.cli.commands.exceptions import StagedArtifactsChangedError from sqlbuild.cli.compile.models import ( + PendingStaticSqlTest, + PlannedStaticSqlTests, SqlTestArtifactCacheRecord, SqlTestArtifactIdentityContext, ) @@ -124,8 +127,9 @@ def write_static_compile_target( adapter: BaseAdapter, project: CompiledProject, manifest: dict[str, object] | None = None, + planned_tests: Callable[[], PlannedStaticSqlTests] | None = None, ) -> WrittenTarget: - """Write offline compiled output files under target_dir.""" + """Write offline compiled output files under target_dir, reusing SQL tests planned ahead.""" remove_stale_files: bool = (target_dir / _COMPILED_DIR).is_dir() target_dir.mkdir(parents=True, exist_ok=True) @@ -152,6 +156,7 @@ def write_static_compile_target( adapter=adapter, project=project, check_existing=remove_stale_files, + planned_tests=planned_tests, ) managed_paths.update(test_paths) if remove_stale_files: @@ -442,27 +447,60 @@ def _write_tests( def static_sql_test_planning_diagnostics( - *, target_dir: Path, adapter: BaseAdapter, project: CompiledProject + *, + target_dir: Path, + adapter: BaseAdapter, + project: CompiledProject, + planned_tests: Callable[[], PlannedStaticSqlTests] | None = None, ) -> tuple[CompilerDiagnostic, ...]: """Plan uncached SQL tests and report their errors without writing artifacts or cache.""" - pending, _, _ = _partition_cached_static_tests( - tests_root=target_dir / _COMPILED_DIR / _TESTS_DIR, adapter=adapter, project=project - ) - tests: tuple[CompiledSqlTest, ...] = tuple(test for test, _, _ in pending) - artifacts: tuple[NativeSqlTestArtifact, ...] = _plan_static_test_artifacts( - adapter=adapter, project=project, tests=tests + planned: PlannedStaticSqlTests = _resolve_planned_static_tests( + target_dir=target_dir, adapter=adapter, project=project, planned_tests=planned_tests ) diagnostics: list[CompilerDiagnostic] = [] - for test, artifact in zip(tests, artifacts, strict=True): - diagnostics.extend(_sql_test_artifact_diagnostics(test=test, artifact=artifact)) + for pending, artifact in zip(planned.pending, planned.artifacts, strict=True): + diagnostics.extend(_sql_test_artifact_diagnostics(test=pending.test, artifact=artifact)) return tuple(diagnostics) +def plan_static_sql_tests( + *, target_dir: Path, adapter: BaseAdapter, project: CompiledProject +) -> PlannedStaticSqlTests: + """Plan every SQL test whose cached artifact under target_dir cannot be reused.""" + + tests_root: Path = target_dir / _COMPILED_DIR / _TESTS_DIR + pending, cached_paths, cached_records = _partition_cached_static_tests( + tests_root=tests_root, adapter=adapter, project=project + ) + return PlannedStaticSqlTests( + tests_root=tests_root, + pending=pending, + artifacts=_plan_static_test_artifacts( + adapter=adapter, project=project, tests=tuple(item.test for item in pending) + ), + cached_paths=frozenset(cached_paths), + cached_records=cached_records, + ) + + +def _resolve_planned_static_tests( + *, + target_dir: Path, + adapter: BaseAdapter, + project: CompiledProject, + planned_tests: Callable[[], PlannedStaticSqlTests] | None, +) -> PlannedStaticSqlTests: + planned: PlannedStaticSqlTests | None = planned_tests() if planned_tests is not None else None + if planned is not None and planned.tests_root == target_dir / _COMPILED_DIR / _TESTS_DIR: + return planned + return plan_static_sql_tests(target_dir=target_dir, adapter=adapter, project=project) + + def _partition_cached_static_tests( *, tests_root: Path, adapter: BaseAdapter, project: CompiledProject ) -> tuple[ - list[tuple[CompiledSqlTest, str | None, str | None]], + tuple[PendingStaticSqlTest, ...], set[Path], dict[str, SqlTestArtifactCacheRecord], ]: @@ -484,7 +522,7 @@ def _partition_cached_static_tests( project=project, tests=project.sql_tests, ) - pending: list[tuple[CompiledSqlTest, str | None, str | None]] = [] + pending: list[PendingStaticSqlTest] = [] for test in project.sql_tests: record_key: str | None = None artifact_identity: str | None = None @@ -506,8 +544,12 @@ def _partition_cached_static_tests( managed_paths.add(cached_path) current_records[record_key] = cached_record continue - pending.append((test, record_key, artifact_identity)) - return pending, managed_paths, current_records + pending.append( + PendingStaticSqlTest( + test=test, record_key=record_key, artifact_identity=artifact_identity + ) + ) + return tuple(pending), managed_paths, current_records def _plan_static_test_artifacts( @@ -534,20 +576,21 @@ def _write_static_tests( adapter: BaseAdapter, project: CompiledProject, check_existing: bool, + planned_tests: Callable[[], PlannedStaticSqlTests] | None, ) -> tuple[set[Path], tuple[CompilerDiagnostic, ...]]: """Write offline SQL-native test SQL and report uncached planning errors.""" diagnostics: list[CompilerDiagnostic] = [] - tests_root: Path = target_dir / _COMPILED_DIR / _TESTS_DIR - pending, managed_paths, current_records = _partition_cached_static_tests( - tests_root=tests_root, adapter=adapter, project=project - ) - native_artifacts: tuple[NativeSqlTestArtifact, ...] = _plan_static_test_artifacts( - adapter=adapter, project=project, tests=tuple(test for test, _, _ in pending) + planned: PlannedStaticSqlTests = _resolve_planned_static_tests( + target_dir=target_dir, adapter=adapter, project=project, planned_tests=planned_tests ) - for (test, record_key, artifact_identity), artifact in zip( - pending, native_artifacts, strict=True - ): + tests_root: Path = planned.tests_root + managed_paths: set[Path] = set(planned.cached_paths) + current_records: dict[str, SqlTestArtifactCacheRecord] = dict(planned.cached_records) + for pending, artifact in zip(planned.pending, planned.artifacts, strict=True): + test: CompiledSqlTest = pending.test + record_key: str | None = pending.record_key + artifact_identity: str | None = pending.artifact_identity test_path: Path = tests_root / compiled_sql_test_output_path( test=test, model_names=artifact.model_names, diff --git a/src/sqlbuild/cli/commands/classes/prepared_compile_artifacts.py b/src/sqlbuild/cli/commands/classes/prepared_compile_artifacts.py index 473b1213ed..9bbab04ad2 100644 --- a/src/sqlbuild/cli/commands/classes/prepared_compile_artifacts.py +++ b/src/sqlbuild/cli/commands/classes/prepared_compile_artifacts.py @@ -1,4 +1,4 @@ -"""Own staged cold compile artifacts and publish them only at the normal write phase.""" +"""Prepare compile artifacts in the background and publish them only at the normal write phase.""" from __future__ import annotations @@ -14,11 +14,14 @@ from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.cli.commands._helpers.compile.target_writer import ( + plan_static_sql_tests, publish_static_compile_target, staged_artifact_files, + static_sql_test_planning_diagnostics, write_static_compile_target, ) from sqlbuild.cli.commands.exceptions import StagedArtifactsChangedError +from sqlbuild.cli.compile.models import PlannedStaticSqlTests from sqlbuild.cli.output.models import WrittenTarget from sqlbuild.compiler.compile.models import CompiledProject, CompilerDiagnostic @@ -32,7 +35,7 @@ class PreparedCompileArtifacts: - """Own one bounded preparation task and its disposable output directory.""" + """Own one background preparation task, either staged artifacts or planned SQL tests only.""" def __init__(self, *, enabled: bool) -> None: self.enabled: bool = enabled @@ -41,6 +44,8 @@ def __init__(self, *, enabled: bool) -> None: self._created_target_dir: Path | None = None self._executor: ThreadPoolExecutor | None = None self._future: Future[_PreparedResult] | None = None + self._planned_tests: Future[PlannedStaticSqlTests] | None = None + self._planning_inputs: tuple[CompiledProject, BaseAdapter, Path] | None = None def __enter__(self) -> PreparedCompileArtifacts: return self @@ -73,16 +78,48 @@ def remove_abandoned_staging(self, *, target_dir: Path) -> None: if not _remove_abandoned_directory(directory): return - def start(self, *, project: CompiledProject, adapter: BaseAdapter, target_dir: Path) -> None: - """Stage artifacts beside target_dir so publication can move them instead of copying.""" + def start( + self, + *, + project: CompiledProject, + adapter: BaseAdapter, + target_dir: Path, + stage_artifacts: bool, + ) -> None: + """Stage artifacts beside target_dir when worthwhile, otherwise plan SQL tests ahead.""" + if not self.enabled: + return if ( - not self.enabled - or project.compile_cache_dir is not None - or len(project.models) < _MIN_PREPARED_ARTIFACT_MODELS - or len(project.models) > _MAX_PREPARED_ARTIFACT_MODELS + stage_artifacts + and _stages_artifacts(project) + and self._start_staging(project=project, adapter=adapter, target_dir=target_dir) ): return + if project.sql_tests: + self._start_test_planning(project=project, adapter=adapter, target_dir=target_dir) + + def _start_test_planning( + self, *, project: CompiledProject, adapter: BaseAdapter, target_dir: Path + ) -> None: + """Plan SQL tests against the published target while the compile continues.""" + + self._planning_inputs = (project, adapter, target_dir) + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-tests") + context: Context = copy_context() + + def plan() -> PlannedStaticSqlTests: + return context.run( + plan_static_sql_tests, target_dir=target_dir, adapter=adapter, project=project + ) + + self._planned_tests = self._executor.submit(plan) + + def _start_staging( + self, *, project: CompiledProject, adapter: BaseAdapter, target_dir: Path + ) -> bool: + """Stage artifacts beside target_dir so publication can move them instead of copying.""" + staging_dir: Path = target_dir / f"{_STAGING_PREFIX}{uuid.uuid4().hex}" try: if not target_dir.is_dir(): @@ -92,7 +129,7 @@ def start(self, *, project: CompiledProject, adapter: BaseAdapter, target_dir: P self._owner_lock.acquire() _create_staging_directory(staging_dir) except OSError: - return + return False self._staging_dir = staging_dir self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-artifacts") context: Context = copy_context() @@ -107,10 +144,20 @@ def prepare() -> _PreparedResult: return written, staged_artifact_files(target_dir=staging_dir) self._future = self._executor.submit(prepare) + return True def publish( self, *, target_dir: Path, manifest: dict[str, object] | None ) -> WrittenTarget | None: + if self._planned_tests is not None and self._planning_inputs is not None: + project, adapter, _ = self._planning_inputs + return write_static_compile_target( + target_dir=target_dir, + adapter=adapter, + project=project, + manifest=manifest, + planned_tests=self._planned_tests.result, + ) if self._future is None or self._staging_dir is None: return None try: @@ -131,8 +178,16 @@ def publish( ) def planning_diagnostics(self) -> tuple[CompilerDiagnostic, ...] | None: - """Return staged SQL test planning diagnostics without publishing any artifact.""" + """Return background SQL test planning diagnostics without publishing any artifact.""" + if self._planned_tests is not None and self._planning_inputs is not None: + project, adapter, target_dir = self._planning_inputs + return static_sql_test_planning_diagnostics( + target_dir=target_dir, + adapter=adapter, + project=project, + planned_tests=self._planned_tests.result, + ) if self._future is None: return None try: @@ -141,6 +196,15 @@ def planning_diagnostics(self) -> tuple[CompilerDiagnostic, ...] | None: return None +def _stages_artifacts(project: CompiledProject) -> bool: + """Staging pays off for uncached, mid-sized projects whose memory stays within budget.""" + + return ( + project.compile_cache_dir is None + and _MIN_PREPARED_ARTIFACT_MODELS <= len(project.models) <= _MAX_PREPARED_ARTIFACT_MODELS + ) + + def _owner_lock_path(staging_dir: Path) -> Path: return staging_dir.with_name(f"{staging_dir.name}{_LOCK_SUFFIX}") diff --git a/src/sqlbuild/cli/compile/models.py b/src/sqlbuild/cli/compile/models.py index a29dfa2f69..ffdf732c95 100644 --- a/src/sqlbuild/cli/compile/models.py +++ b/src/sqlbuild/cli/compile/models.py @@ -8,10 +8,15 @@ from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.cli.commands.types import CompileLineageMode from sqlbuild.cli.output.models import WrittenTarget -from sqlbuild.compiler.compile.models import CompiledObjectKey, CompilerDiagnostic +from sqlbuild.compiler.compile.models import ( + CompiledObjectKey, + CompiledSqlTest, + CompilerDiagnostic, +) from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs from sqlbuild.compiler.lineage.models import ProjectColumnLineage from sqlbuild.compiler.pipeline.models import ProjectGraph +from sqlbuild.compiler.planner.models import NativeSqlTestArtifact @dataclass(frozen=True) @@ -90,3 +95,23 @@ class SqlTestArtifactIdentityContext: common_identity: str model_identities: dict[str, str] + + +@dataclass(frozen=True) +class PendingStaticSqlTest: + """One SQL test that needs a freshly planned artifact, with its optional cache identity.""" + + test: CompiledSqlTest + record_key: str | None + artifact_identity: str | None + + +@dataclass(frozen=True) +class PlannedStaticSqlTests: + """SQL-test artifacts planned for one target directory before any file is written.""" + + tests_root: Path + pending: tuple[PendingStaticSqlTest, ...] + artifacts: tuple[NativeSqlTestArtifact, ...] + cached_paths: frozenset[Path] + cached_records: dict[str, SqlTestArtifactCacheRecord] 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 51001c35ed..71b10cd215 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py @@ -31,6 +31,25 @@ class PreparedArtifactsCompileTestCase: expected_exit_code: int = 0 +@dataclass(frozen=True) +class BackgroundTestPlanningCompileTestCase: + description: str + compile_args: tuple[str, ...] + rules_config: str + prepare_project: Callable[..., None] + max_prepared_models: int = 5000 + expected_background_plans: tuple[str, ...] = ("sqlbuild-tests",) + + +@dataclass(frozen=True) +class BackgroundTestPlanningFailureTestCase: + description: str + compile_args: tuple[str, ...] + test_options: str + patch_native_planner: Callable[..., None] + expected_error_fragment: str + + @dataclass(frozen=True) class AbandonedStagingCompileTestCase: description: str diff --git a/tests/integration/src/sqlbuild/cli/commands/main/helpers.py b/tests/integration/src/sqlbuild/cli/commands/main/helpers.py index c2d82be487..9e2971265c 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/helpers.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/helpers.py @@ -3,6 +3,7 @@ import json import os import tempfile +import threading from collections.abc import Callable from pathlib import Path from typing import Any, cast @@ -17,12 +18,18 @@ ) from sqlbuild.adapter.contract.types import RetentionChangePhase from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter +from sqlbuild.cli.commands._helpers.compile import target_writer +from sqlbuild.cli.commands.classes import prepared_compile_artifacts +from sqlbuild.cli.commands.classes.prepared_compile_artifacts import PreparedCompileArtifacts from sqlbuild.cli.commands.main.entrypoint.entry import main +from sqlbuild.cli.compile.models import PlannedStaticSqlTests from sqlbuild.compiler.compile._helpers.assembly import source_bindings as source_bindings_module from sqlbuild.compiler.compile.models import PolyglotAnalysisResult from sqlbuild.compiler.discovery._helpers.filesystem import ( model_files as discovery_model_files_module, ) +from sqlbuild.compiler.planner.exceptions import NativeSqlTestPlanningError +from sqlbuild.compiler.planner.models import NativeSqlTestArtifact from sqlbuild.spec.contracts.models import SourceLocation @@ -59,6 +66,127 @@ def write_prepared_artifacts_project(*, project_dir: Path, model_count: int) -> ) +SELECTED_RULES_CONFIG: str = '[rules]\nselect = ["SQBRSQL021"]\n' +TEST_PLANNING_MODELS: int = 6 + + +def write_test_planning_project( + *, project_dir: Path, model_count: int, rules_config: str, test_options: str = "" +) -> None: + """Write a chain of DuckDB models with one SQL test per downstream model.""" + + project_dir.mkdir(parents=True, exist_ok=True) + (project_dir / "sqlbuild_project.toml").write_text( + f'name = "orders"\nadapter = "duckdb"\n{rules_config}', encoding="utf-8" + ) + models: Path = project_dir / "models" + tests: Path = project_dir / "tests" / "unit" + models.mkdir() + tests.mkdir(parents=True) + header: str = ( + "MODEL (description 'Test model.', materialized table, contract enforced, " + "columns (order_id (type INTEGER)));\n" + ) + (models / "orders_000.sql").write_text( + header + "SELECT CAST(0 AS INTEGER) AS order_id\n", encoding="utf-8" + ) + for index in range(1, model_count): + upstream: str = f"orders_{index - 1:03}" + (models / f"orders_{index:03}.sql").write_text( + header + f'SELECT CAST(order_id + 1 AS INTEGER) AS order_id FROM __ref("{upstream}")\n', + encoding="utf-8", + ) + write_order_test( + tests_dir=tests, index=index, expected_order_id=2, test_options=test_options + ) + + +def write_order_test( + *, tests_dir: Path, index: int, expected_order_id: int, test_options: str = "" +) -> None: + """Write the SQL test for one chained order model with a mocked upstream model.""" + + (tests_dir / f"orders_{index:03}.sql").write_text( + f'TEST (name "orders_{index:03}_case"{test_options});\n' + "WITH\n" + f"__ref__orders_{index - 1:03} AS (SELECT CAST(1 AS INTEGER) AS order_id),\n" + f"__expected__orders_{index:03} AS " + f"(SELECT CAST({expected_order_id} AS INTEGER) AS order_id)\n" + "SELECT 1\n", + encoding="utf-8", + ) + + +def compiled_files(project_dir: Path) -> dict[str, bytes]: + """Return every compiled artifact under target/compiled keyed by relative path.""" + + compiled: Path = project_dir / "target" / "compiled" + return { + str(path.relative_to(compiled)): path.read_bytes() + for path in sorted(compiled.rglob("*.sql")) + } + + +def semantic_compile_payload(output: str) -> dict[str, object]: + """Return compile JSON without its timings and version, which vary between runs.""" + + payload: dict[str, object] = json.loads(output) + payload.pop("compile_timings", None) + payload.pop("version", None) + return payload + + +def compile_project_json( + *, project_dir: Path, compile_args: tuple[str, ...], capsys: pytest.CaptureFixture[str] +) -> tuple[int, dict[str, object]]: + """Compile through the CLI and return the exit code with the semantic JSON payload.""" + + exit_code: int = main(["--project-dir", str(project_dir), "compile", "--json", *compile_args]) + return exit_code, semantic_compile_payload(capsys.readouterr().out) + + +def keep_project_unchanged(*, project_dir: Path, capsys: pytest.CaptureFixture[str]) -> None: + """Leave a freshly written project cold, with no target directory.""" + + +def compile_then_edit_one_test(*, project_dir: Path, capsys: pytest.CaptureFixture[str]) -> None: + """Warm the compile cache, then change one SQL test so only it must be planned again.""" + + assert compile_project_json(project_dir=project_dir, compile_args=(), capsys=capsys)[0] == 0 + write_order_test(tests_dir=project_dir / "tests" / "unit", index=2, expected_order_id=3) + + +def keep_native_planner(monkeypatch: pytest.MonkeyPatch) -> None: + """Leave the real native SQL-test planner in place.""" + + +def fail_native_planner(monkeypatch: pytest.MonkeyPatch) -> None: + """Make the native SQL-test planner fail as an unexpected internal error would.""" + + def fail(**kwargs: object) -> tuple[NativeSqlTestArtifact, ...]: + raise NativeSqlTestPlanningError("native SQL-test planning failed: worker stopped") + + monkeypatch.setattr(target_writer, "plan_and_render_sql_test_artifacts", fail) + + +def record_test_planning_threads(monkeypatch: pytest.MonkeyPatch) -> list[str]: + """Record the thread-name prefix of every background SQL-test planning task.""" + + threads: list[str] = [] + plan: Callable[..., PlannedStaticSqlTests] = prepared_compile_artifacts.plan_static_sql_tests + + def recorded_plan(**kwargs: Any) -> PlannedStaticSqlTests: + threads.append(threading.current_thread().name.rpartition("_")[0]) + return plan(**kwargs) + + monkeypatch.setattr(prepared_compile_artifacts, "plan_static_sql_tests", recorded_plan) + return threads + + +def skip_artifact_preparation(self: PreparedCompileArtifacts, **kwargs: object) -> None: + """Disable background preparation so the write phase plans SQL tests inline.""" + + def has_second_filesystem(root: Path) -> bool: """Return whether root is a writable directory on another device than the temp directory.""" diff --git a/tests/integration/src/sqlbuild/cli/commands/main/test_compile_prepared_artifacts.py b/tests/integration/src/sqlbuild/cli/commands/main/test_compile_prepared_artifacts.py index 20374275d6..8d98c44dd4 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/test_compile_prepared_artifacts.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/test_compile_prepared_artifacts.py @@ -2,6 +2,7 @@ import json import os +import shutil from contextlib import ExitStack from pathlib import Path from tempfile import TemporaryDirectory @@ -11,23 +12,168 @@ from filelock import FileLock from sqlbuild.cli.commands.classes import prepared_compile_artifacts +from sqlbuild.cli.commands.classes.prepared_compile_artifacts import PreparedCompileArtifacts from sqlbuild.cli.commands.main.entrypoint.entry import main from tests.integration.src.sqlbuild.cli.commands.main._test_types import ( AbandonedStagingCompileTestCase, + BackgroundTestPlanningCompileTestCase, + BackgroundTestPlanningFailureTestCase, CrossDeviceArtifactsCompileTestCase, PreparedArtifactsCompileTestCase, ) from tests.integration.src.sqlbuild.cli.commands.main.helpers import ( + SELECTED_RULES_CONFIG, + TEST_PLANNING_MODELS, + compile_project_json, + compile_then_edit_one_test, + compiled_files, + fail_native_planner, has_second_filesystem, + keep_native_planner, + keep_project_unchanged, + record_test_planning_threads, + skip_artifact_preparation, staging_directories, unavailable_staging_directory, write_prepared_artifacts_project, write_staging_directory, + write_test_planning_project, ) _SECOND_FILESYSTEM: Path = Path("/dev/shm") +@pytest.mark.parametrize( + "test_case", + ( + BackgroundTestPlanningCompileTestCase( + "cold compile with the compile cache", + (), + SELECTED_RULES_CONFIG, + keep_project_unchanged, + ), + BackgroundTestPlanningCompileTestCase( + "warm compile after one SQL test edit", + (), + SELECTED_RULES_CONFIG, + compile_then_edit_one_test, + ), + BackgroundTestPlanningCompileTestCase( + "no-cache compile above the staged artifact limit", + ("--no-cache",), + SELECTED_RULES_CONFIG, + keep_project_unchanged, + max_prepared_models=1, + ), + BackgroundTestPlanningCompileTestCase( + "no-cache compile without selected rules", ("--no-cache",), "", keep_project_unchanged + ), + BackgroundTestPlanningCompileTestCase( + "no-cache compile that stages every artifact", + ("--no-cache",), + SELECTED_RULES_CONFIG, + keep_project_unchanged, + expected_background_plans=(), + ), + ), + ids=lambda case: case.description, +) +def test_given_compile_mode_when_preparing_artifacts_in_background_then_output_matches_inline( + test_case: BackgroundTestPlanningCompileTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(prepared_compile_artifacts, "_MIN_PREPARED_ARTIFACT_MODELS", 1) + monkeypatch.setattr( + prepared_compile_artifacts, "_MAX_PREPARED_ARTIFACT_MODELS", test_case.max_prepared_models + ) + background: Path = tmp_path / "background" + write_test_planning_project( + project_dir=background, + model_count=TEST_PLANNING_MODELS, + rules_config=test_case.rules_config, + ) + test_case.prepare_project(project_dir=background, capsys=capsys) + inline: Path = tmp_path / "inline" + shutil.copytree(background, inline) + planning_threads: list[str] = record_test_planning_threads(monkeypatch) + + background_result: tuple[int, dict[str, object]] = compile_project_json( + project_dir=background, compile_args=test_case.compile_args, capsys=capsys + ) + with monkeypatch.context() as inline_patch: + inline_patch.setattr(PreparedCompileArtifacts, "start", skip_artifact_preparation) + inline_result: tuple[int, dict[str, object]] = compile_project_json( + project_dir=inline, compile_args=test_case.compile_args, capsys=capsys + ) + + assert tuple(planning_threads) == test_case.expected_background_plans + assert background_result[0] == 0 + assert background_result == inline_result + assert background_result[1]["has_errors"] is False + artifacts: dict[str, bytes] = compiled_files(background) + assert artifacts == compiled_files(inline) + assert sum(path.startswith("tests/") for path in artifacts) == TEST_PLANNING_MODELS - 1 + assert staging_directories(background / "target") == [] + + +@pytest.mark.parametrize( + "test_case", + ( + BackgroundTestPlanningFailureTestCase( + "planner input error with the compile cache", + (), + ', cursor_start "2026-02-01"', + keep_native_planner, + "declares cursor_start or cursor_end", + ), + BackgroundTestPlanningFailureTestCase( + "unexpected native planning failure without the compile cache", + ("--no-cache",), + "", + fail_native_planner, + "worker stopped", + ), + ), + ids=lambda case: case.description, +) +def test_given_background_planning_fails_when_compiling_then_error_matches_inline_planning( + test_case: BackgroundTestPlanningFailureTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + test_case.patch_native_planner(monkeypatch) + background: Path = tmp_path / "background" + write_test_planning_project( + project_dir=background, + model_count=TEST_PLANNING_MODELS, + rules_config=SELECTED_RULES_CONFIG, + test_options=test_case.test_options, + ) + inline: Path = tmp_path / "inline" + shutil.copytree(background, inline) + planning_threads: list[str] = record_test_planning_threads(monkeypatch) + + background_result: tuple[int, dict[str, object]] = compile_project_json( + project_dir=background, compile_args=test_case.compile_args, capsys=capsys + ) + with monkeypatch.context() as inline_patch: + inline_patch.setattr(PreparedCompileArtifacts, "start", skip_artifact_preparation) + inline_result: tuple[int, dict[str, object]] = compile_project_json( + project_dir=inline, compile_args=test_case.compile_args, capsys=capsys + ) + + assert planning_threads == ["sqlbuild-tests"] + assert background_result[0] == 1 + assert background_result == inline_result + assert test_case.expected_error_fragment in json.dumps(background_result[1]) + artifacts: dict[str, bytes] = compiled_files(background) + assert artifacts == compiled_files(inline) + assert not any(path.startswith("tests/") for path in artifacts) + + @pytest.mark.parametrize( "test_case", ( diff --git a/tests/unit/src/sqlbuild/cli/commands/main/compile/test_compile_command.py b/tests/unit/src/sqlbuild/cli/commands/main/compile/test_compile_command.py index 01b7ed6c4a..9a643a83dc 100644 --- a/tests/unit/src/sqlbuild/cli/commands/main/compile/test_compile_command.py +++ b/tests/unit/src/sqlbuild/cli/commands/main/compile/test_compile_command.py @@ -551,12 +551,12 @@ def test_given_lineage_disabled_when_building_compile_lineage_then_skips_analyze "test_case", [ RuleGatedTestPlanningTestCase( - description="inline writer plans once for withheld artifacts", + description="background test planning is reused for withheld artifacts", min_prepared_models=128, test_header='name "orders_case"', expected_codes=("SQBRMODEL102", "S000"), expected_message_fragment="references __source('raw_refunds') which has no mock", - expected_inline_planning_calls=1, + expected_inline_planning_calls=0, ), RuleGatedTestPlanningTestCase( description="background preparation reuses staged planning", @@ -567,12 +567,12 @@ def test_given_lineage_disabled_when_building_compile_lineage_then_skips_analyze expected_inline_planning_calls=0, ), RuleGatedTestPlanningTestCase( - description="inline planning input error becomes a diagnostic", + description="background test planning input error becomes a diagnostic", min_prepared_models=128, test_header='name "orders_case", cursor_start "2026-02-01"', expected_codes=("SQBRMODEL102", "P001"), expected_message_fragment="declares cursor_start or cursor_end", - expected_inline_planning_calls=1, + expected_inline_planning_calls=0, ), RuleGatedTestPlanningTestCase( description="background planning input error becomes a diagnostic", diff --git a/tests/unit/src/sqlbuild/cli/commands/main/compile/test_staged_artifact_publication.py b/tests/unit/src/sqlbuild/cli/commands/main/compile/test_staged_artifact_publication.py index adffe372eb..1cc116e81b 100644 --- a/tests/unit/src/sqlbuild/cli/commands/main/compile/test_staged_artifact_publication.py +++ b/tests/unit/src/sqlbuild/cli/commands/main/compile/test_staged_artifact_publication.py @@ -172,6 +172,7 @@ def remove_while_locked(path: Any, *args: Any, **kwargs: Any) -> None: project=build_static_target_writer_project(), adapter=DuckDbAdapter(), target_dir=target, + stage_artifacts=True, ) written: WrittenTarget | None = artifacts.publish(target_dir=target, manifest=None) @@ -208,6 +209,7 @@ def test_given_staging_owner_lost_when_publishing_then_compile_fails_without_del project=build_static_target_writer_project(), adapter=DuckDbAdapter(), target_dir=target, + stage_artifacts=True, ) _ = artifacts.planning_diagnostics() for relative in test_case.removed_staged_files: