From cb42722b9027394acff270cc9b86bf647b615c76 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Tue, 6 Oct 2026 05:22:27 +0100 Subject: [PATCH 1/4] perf(rules): start custom-rule hosts early and keep their start-up cheap Hosts now start before the parent writes the project payload and wait for it to be published, so interpreter start-up overlaps the payload encoding. A failed write publishes an empty payload so waiting hosts exit promptly. Hosts load the project with cyclic collection paused and then freeze it, and the rule-test harness is imported only when a test calls it, so rule modules no longer pull the compile pipeline into every host. --- .../rule_engine/_helpers/host/custom_host.py | 84 ++++++-- .../_helpers/host/custom_host_pool.py | 100 +++++++-- src/sqlbuild/rule_engine/constants.py | 3 + .../rule_engine/main/evaluate_rule.py | 5 +- .../rule_engine/_helpers/host/_test_types.py | 32 +++ .../rule_engine/_helpers/host/helpers.py | 49 +++++ .../_helpers/host/test_custom_host_pool.py | 191 +++++++++++++++++- .../main/evaluate_rule/_test_types.py | 8 + .../evaluate_rule/test_authoring_imports.py | 55 +++++ 9 files changed, 487 insertions(+), 40 deletions(-) create mode 100644 tests/unit/src/sqlbuild/rule_engine/main/evaluate_rule/test_authoring_imports.py diff --git a/src/sqlbuild/rule_engine/_helpers/host/custom_host.py b/src/sqlbuild/rule_engine/_helpers/host/custom_host.py index f19d2b5347..b6223e5a55 100644 --- a/src/sqlbuild/rule_engine/_helpers/host/custom_host.py +++ b/src/sqlbuild/rule_engine/_helpers/host/custom_host.py @@ -3,14 +3,16 @@ 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 -from typing import Any +from typing import Any, BinaryIO from sqlbuild.compiler.compile.models import CompiledProject from sqlbuild.rule_engine._helpers.engine.catalogue import build_catalogue @@ -20,8 +22,11 @@ from sqlbuild.rule_engine.constants import ( CUSTOM_HOST_CANCELLED_EXIT_CODE, CUSTOM_HOST_CANCELLED_MESSAGE, + 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_LAUNCH_MODULE, CUSTOM_HOST_MAX_TRACKED_READS, CUSTOM_HOST_PROTOCOL_VERSION, @@ -61,7 +66,7 @@ def main() -> int: cancel_marker: Path | None = _cancel_marker(payload) if cancel_marker is not None and cancel_marker.exists(): raise HostCancelledError(CUSTOM_HOST_CANCELLED_MESSAGE) - project, config = _decode_inputs(payload) + project, config = _decode_inputs(payload=payload, cancel_marker=cancel_marker) project_dir: Path = Path(str(payload["project_dir"])).resolve() dialect: str = str(payload.get("dialect", "generic")) verify_determinism: bool = payload.get("verify_determinism") is True @@ -111,18 +116,27 @@ def main() -> int: return 0 -def _decode_inputs(payload: dict[str, Any]) -> tuple[CompiledProject, RulesConfig]: +def _decode_inputs( + *, payload: dict[str, Any], cancel_marker: Path | None +) -> tuple[CompiledProject, 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"])) - if ( - not input_path.is_absolute() - or input_path.is_symlink() - or not input_path.resolve().is_relative_to(input_root) - ): + 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( + path=input_path, + cancel_markers=tuple( + marker + for marker in (cancel_marker, _input_marker(payload=payload, key="abandoned_marker")) + if marker is not None + ), + parent_pid=_expected_parent(payload), + ) + 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 = pickle.load(handle) + decoded: object = _load_long_lived(handle) if ( not isinstance(decoded, tuple) or len(decoded) != CUSTOM_HOST_INPUT_TUPLE_SIZE @@ -133,20 +147,62 @@ def _decode_inputs(payload: dict[str, Any]) -> tuple[CompiledProject, RulesConfi return decoded -def _cancel_marker(payload: dict[str, Any]) -> Path | None: - """Return the run's cancellation marker, which must live beside the project payload.""" +def await_inputs(*, path: Path, cancel_markers: tuple[Path, ...], parent_pid: int) -> None: + """Wait for the payload the parent publishes after starting this host, while it still runs.""" + + deadline: float = time.monotonic() + CUSTOM_HOST_INPUT_WAIT_SECONDS + while not path.exists(): + if any(marker.exists() for marker in cancel_markers): + raise HostCancelledError(CUSTOM_HOST_CANCELLED_MESSAGE) + if os.getppid() != parent_pid: + raise RulesError("custom host parent exited before writing the project payload") + if time.monotonic() > deadline: + raise RulesError("custom host project payload was never written") + time.sleep(CUSTOM_HOST_INPUT_POLL_SECONDS) + + +def _expected_parent(payload: dict[str, Any]) -> int: + """Return the parent that started this host; a relaying launcher stands in for it.""" + + value: object = payload.get("parent_pid") + if os.name == CUSTOM_HOST_EXEC_OS_NAME and isinstance(value, int) and value > 0: + return value + return os.getppid() + - value: object = payload.get("cancel_marker") +def _input_marker(*, payload: dict[str, Any], key: str) -> Path | None: + """Return a signal file the run placed beside the project payload, if it named one.""" + + value: object = payload.get(key) if value is None: return None project_dir: Path = Path(str(payload["project_dir"])).resolve() input_root: Path = (project_dir / "target" / "rules-cache" / "host-inputs").resolve() marker: Path = Path(str(value)) - if not marker.is_absolute() or not marker.resolve().is_relative_to(input_root): - raise RulesError("custom host cancellation marker path is invalid") + if not marker.is_absolute() or not marker.parent.resolve().is_relative_to(input_root): + raise RulesError(f"custom host {key} path is invalid") 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.""" + + return _input_marker(payload=payload, key="cancel_marker") + + def _decode_plan(value: object) -> CustomRulePlan: if not isinstance(value, dict): raise RulesError("custom host plan must be an object") diff --git a/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py b/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py index e84492b29a..8d5710be3a 100644 --- a/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py +++ b/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py @@ -2,10 +2,13 @@ from __future__ import annotations +import os import pickle import shutil import sys import tempfile +import uuid +from collections.abc import Callable, Iterable from concurrent.futures import FIRST_EXCEPTION, Future, ThreadPoolExecutor, wait from dataclasses import replace from pathlib import Path @@ -17,6 +20,7 @@ from sqlbuild.compiler.compile.models import CompiledProject from sqlbuild.rule_engine._helpers.host.host_capacity import available_cores from sqlbuild.rule_engine.constants import ( + CUSTOM_HOST_ABANDONED_SUFFIX, CUSTOM_HOST_CANCELLED_MARKER, CUSTOM_HOST_LAUNCH_MODULE, CUSTOM_HOST_MAX_TRACKED_READS, @@ -64,43 +68,78 @@ def run_custom_hosts( ) input_dir: Path = project_dir / "target" / "rules-cache" / "host-inputs" input_dir.mkdir(parents=True, exist_ok=True) - input_path: Path | None = None + input_path: Path = input_dir / f"project-{uuid.uuid4().hex}.pickle" + abandoned: Path = input_path.with_suffix(CUSTOM_HOST_ABANDONED_SUFFIX) + payload: dict[str, object] = { + "project_pickle_path": str(input_path.resolve()), + "abandoned_marker": str(abandoned.resolve()), + "parent_pid": os.getpid(), + "project_dir": str(project_dir.resolve()), + "dialect": dialect, + "verify_determinism": verify_determinism, + "track_reads": track_reads, + } try: - with tempfile.NamedTemporaryFile( - dir=input_dir, prefix="project-", suffix=".pickle", delete=False - ) as handle: - input_path = Path(handle.name) - pickle.dump((project, config), handle) - payload: dict[str, object] = { - "project_pickle_path": str(input_path.resolve()), - "project_dir": str(project_dir.resolve()), - "dialect": dialect, - "verify_determinism": verify_determinism, - "track_reads": track_reads, - } - responses: list[CustomHostSlice] = ( - [_run_host(payload={**payload, "plan": plan})] - if len(plans) == 1 - else _run_partitions(payload=payload, plan=plan, plans=plans, input_dir=input_dir) + responses: list[CustomHostSlice] = _run_partitions( + payload=payload, + plan=plan, + plans=plans, + input_dir=input_dir, + write_inputs=lambda: _write_inputs(path=input_path, project=project, config=config), ) finally: - if input_path is not None: - input_path.unlink(missing_ok=True) + input_path.unlink(missing_ok=True) + abandoned.unlink(missing_ok=True) return _merge( responses=responses, order={invocation: index for index, invocation in enumerate(invocations)}, ) +def _write_inputs(*, path: Path, project: CompiledProject, config: RulesConfig) -> None: + """Publish the project payload atomically; hosts started earlier wait for it to appear.""" + + with tempfile.NamedTemporaryFile( + dir=path.parent, prefix="writing-", suffix=".pickle", delete=False + ) as handle: + partial: Path = Path(handle.name) + try: + pickle.dump((project, config), handle) + except BaseException: + handle.close() + partial.unlink(missing_ok=True) + raise + partial.replace(path) + + +def _abandon_inputs(marker: Path) -> bool: + """Tell waiting hosts no payload will come; report whether they can see the signal.""" + + try: + marker.parent.mkdir(parents=True, exist_ok=True) + marker.touch() + except OSError: + return False + return True + + def _run_partitions( *, payload: dict[str, object], plan: CustomHostPlan, plans: tuple[CustomHostPlan, ...], input_dir: Path, + write_inputs: Callable[[], None], ) -> list[CustomHostSlice]: - """Run partitions; failures and rules with module state fall back to one host.""" + """Start hosts while the payload is written; failures and module state fall back to one host.""" + if len(plans) == 1: + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-custom-host") as pool: + single: Future[CustomHostSlice] = pool.submit( + _run_host, payload={**payload, "plan": plan} + ) + _publish(write_inputs=write_inputs, payload=payload, futures=(single,)) + return [single.result()] cancel_dir: Path = Path(tempfile.mkdtemp(dir=input_dir, prefix="hosts-")) marker: Path = cancel_dir / CUSTOM_HOST_CANCELLED_MARKER futures: list[Future[CustomHostSlice]] = [] @@ -120,6 +159,7 @@ def _run_partitions( ) for partition in plans ] + _publish(write_inputs=write_inputs, payload=payload, futures=futures) done: set[Future[CustomHostSlice]] = wait(futures, return_when=FIRST_EXCEPTION).done failed = any(future.exception() is not None for future in done) if failed: @@ -132,13 +172,29 @@ def _run_partitions( stateful: frozenset[str] = frozenset().union(*(item.stateful for item in slices)) if not stateful: return slices - single: CustomHostSlice = _run_host( + rerun: CustomHostSlice = _run_host( payload={ **payload, "plan": {code: plan[code] if code in stateful else [] for code in plan}, }, ) - return [*(_without_codes(response=item, codes=stateful) for item in slices), single] + return [*(_without_codes(response=item, codes=stateful) for item in slices), rerun] + + +def _publish( + *, + write_inputs: Callable[[], None], + payload: dict[str, object], + futures: Iterable[Future[CustomHostSlice]], +) -> None: + """Write the payload or release hosts; if release fails, the error waits for their deadline.""" + + try: + write_inputs() + except BaseException: + if _abandon_inputs(Path(str(payload["abandoned_marker"]))): + _ = wait(tuple(futures)) + raise def _cancel_hosts(marker: Path) -> None: diff --git a/src/sqlbuild/rule_engine/constants.py b/src/sqlbuild/rule_engine/constants.py index 306024b046..d8ba681b11 100644 --- a/src/sqlbuild/rule_engine/constants.py +++ b/src/sqlbuild/rule_engine/constants.py @@ -16,6 +16,9 @@ NATIVE_RULES_MEMO_VERSION: str = "native-rules-response-v2" NATIVE_RULES_MEMO_FILE: str = "target/rules-cache/bulk/native-response.json" CUSTOM_HOST_INPUT_TUPLE_SIZE: int = 2 +CUSTOM_HOST_ABANDONED_SUFFIX: str = ".abandoned" +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" CUSTOM_HOST_CANCELLED_MARKER: str = "cancelled" CUSTOM_HOST_CANCELLED_EXIT_CODE: int = 75 diff --git a/src/sqlbuild/rule_engine/main/evaluate_rule.py b/src/sqlbuild/rule_engine/main/evaluate_rule.py index 5215487469..d728597f83 100644 --- a/src/sqlbuild/rule_engine/main/evaluate_rule.py +++ b/src/sqlbuild/rule_engine/main/evaluate_rule.py @@ -1,11 +1,12 @@ """Public custom rule test harness.""" -from sqlbuild.rule_engine._helpers.engine.rule_harness import run_rule_case from sqlbuild.rule_engine.models import Rule, RuleCase, RuleResult from sqlbuild.rule_engine.types import RuleCheck def evaluate_rule(*, rule: RuleCheck | Rule, test_case: RuleCase) -> RuleResult: - """Evaluate one rule through discovery, compilation, and rule evaluation.""" + """Evaluate one rule through discovery, compilation, and rule evaluation, loaded on demand.""" + + from sqlbuild.rule_engine._helpers.engine.rule_harness import run_rule_case return run_rule_case(rule=rule, test_case=test_case) diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py index 516e7d4943..efadbee60c 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py @@ -88,3 +88,35 @@ class StateDetectorCostCase: rule_count: int model_count: int expected_min_speedup: float + + +@dataclass(frozen=True) +class HostInputsTestCase: + """Hosts started before their project payload is published, or abandoned without one.""" + + description: str + model_count: int + hosts: int + write_delay_seconds: float + expected_max_seconds: float + + +@dataclass(frozen=True) +class AwaitInputsTestCase: + """One wait for a host payload with a given parent, cancel signal, and payload state.""" + + description: str + reported_parent_pid: int + existing_files: tuple[str, ...] + expected_error: type[Exception] + expected_max_seconds: float + + +@dataclass(frozen=True) +class AwaitPublishedInputsTestCase: + """One wait for a payload that is already published while the parent still runs.""" + + description: str + reported_parent_pid: int + existing_files: tuple[str, ...] + expected_max_seconds: float diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py index 205e102db6..41e1a4a1f7 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py @@ -2,6 +2,7 @@ import os import random +import shutil import time from collections.abc import Callable from dataclasses import replace @@ -429,3 +430,51 @@ def timed_state_history( rules=rules, model_count=model_count, detectors=(detector(rules=rules, root=root),) ) return time.process_time() - started, history + + +def delay_inputs(*, monkeypatch: pytest.MonkeyPatch, delay_seconds: float) -> None: + """Publish every host project payload only after a delay.""" + + write: Callable[..., None] = custom_host_pool._write_inputs + + def delayed_write(*, path: Path, project: CompiledProject, config: RulesConfig) -> None: + time.sleep(delay_seconds) + write(path=path, project=project, config=config) + + monkeypatch.setattr(custom_host_pool, "_write_inputs", delayed_write) + + +def fail_inputs(*, monkeypatch: pytest.MonkeyPatch, delay_seconds: float) -> None: + """Fail to publish every host project payload after a delay.""" + + def failed_write(*, path: Path, project: CompiledProject, config: RulesConfig) -> None: + del path, project, config + time.sleep(delay_seconds) + raise OSError("project payload disk is full") + + monkeypatch.setattr(custom_host_pool, "_write_inputs", failed_write) + + +def fail_inputs_without_directory( + *, monkeypatch: pytest.MonkeyPatch, delay_seconds: float +) -> list[float]: + """Remove the host-input directory, then fail to publish; return when each write failed.""" + + failures: list[float] = [] + + def failed_write(*, path: Path, project: CompiledProject, config: RulesConfig) -> None: + del project, config + time.sleep(delay_seconds) + shutil.rmtree(path.parent) + failures.append(time.monotonic()) + raise OSError("project payload disk is full") + + monkeypatch.setattr(custom_host_pool, "_write_inputs", failed_write) + return failures + + +def touch_files(*, root: Path, names: tuple[str, ...]) -> None: + """Create empty files beneath one folder.""" + + for name in names: + (root / name).touch() 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 7f1eccf405..819d317d9d 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 @@ -6,20 +6,27 @@ import pytest from sqlbuild.compiler.compile.models import CompiledProject -from sqlbuild.rule_engine._helpers.host import custom_host_pool +from sqlbuild.rule_engine._helpers.host import custom_host, custom_host_pool from sqlbuild.rule_engine.constants import CUSTOM_RULES_CACHE_FILE -from sqlbuild.rule_engine.exceptions import RulesError +from sqlbuild.rule_engine.exceptions import HostCancelledError, RulesError from sqlbuild.rule_engine.models import CustomRulesOutcome from tests.unit.src.sqlbuild.rule_engine._helpers.host._test_types import ( + AwaitInputsTestCase, + AwaitPublishedInputsTestCase, HostCancellationTestCase, HostFailureTestCase, + HostInputsTestCase, HostPartitionTestCase, HostPlanTestCase, ) from tests.unit.src.sqlbuild.rule_engine._helpers.host.helpers import ( + delay_inputs, evaluate_order_rules, + fail_inputs, + fail_inputs_without_directory, orders_project, record_signals, + touch_files, use_hosts, write_order_rules, ) @@ -280,5 +287,185 @@ def test_given_cancelled_hosts_when_failing_then_no_process_is_signalled( assert test_case.expected_error_fragment in str(raised.value) +@pytest.mark.parametrize( + "test_case", + ( + HostInputsTestCase( + description="one host waits for a slow payload", + model_count=4, + hosts=1, + write_delay_seconds=1.0, + expected_max_seconds=60.0, + ), + HostInputsTestCase( + description="split hosts wait for a slow payload", + model_count=7, + hosts=3, + write_delay_seconds=1.0, + expected_max_seconds=60.0, + ), + ), + ids=lambda case: case.description, +) +def test_given_hosts_started_before_payload_when_payload_arrives_then_findings_are_unchanged( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostInputsTestCase +) -> None: + write_order_rules(root=tmp_path) + project: CompiledProject = orders_project(model_count=test_case.model_count) + oracle: CustomRulesOutcome = evaluate_order_rules( + project=project, root=tmp_path, cache_enabled=False + ) + launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) + delay_inputs(monkeypatch=monkeypatch, delay_seconds=test_case.write_delay_seconds) + started: float = time.monotonic() + + delayed: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) + + assert time.monotonic() - started < test_case.expected_max_seconds + assert len(launches) == test_case.hosts + assert delayed.findings == oracle.findings + assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) + + +@pytest.mark.parametrize( + "test_case", + ( + HostInputsTestCase( + description="one waiting host is released", + model_count=4, + hosts=1, + write_delay_seconds=0.5, + expected_max_seconds=30.0, + ), + HostInputsTestCase( + description="split waiting hosts are released", + model_count=7, + hosts=3, + write_delay_seconds=0.5, + expected_max_seconds=30.0, + ), + ), + ids=lambda case: case.description, +) +def test_given_payload_write_failure_when_hosts_wait_then_original_error_is_raised_promptly( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostInputsTestCase +) -> None: + write_order_rules(root=tmp_path) + project: CompiledProject = orders_project(model_count=test_case.model_count) + launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) + fail_inputs(monkeypatch=monkeypatch, delay_seconds=test_case.write_delay_seconds) + started: float = time.monotonic() + + with pytest.raises(OSError, match="project payload disk is full"): + _ = evaluate_order_rules(project=project, root=tmp_path) + + assert time.monotonic() - started < test_case.expected_max_seconds + assert len(launches) == test_case.hosts + assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) + + +@pytest.mark.parametrize( + "test_case", + ( + HostInputsTestCase( + description="one host released after its input folder vanished", + model_count=4, + hosts=1, + write_delay_seconds=2.0, + expected_max_seconds=1.0, + ), + HostInputsTestCase( + description="split hosts released after their input folder vanished", + model_count=7, + hosts=3, + write_delay_seconds=2.0, + expected_max_seconds=1.0, + ), + ), + ids=lambda case: case.description, +) +def test_given_input_folder_removed_when_payload_write_fails_then_original_error_returns_promptly( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostInputsTestCase +) -> None: + write_order_rules(root=tmp_path) + project: CompiledProject = orders_project(model_count=test_case.model_count) + launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) + failures: list[float] = fail_inputs_without_directory( + monkeypatch=monkeypatch, delay_seconds=test_case.write_delay_seconds + ) + + with pytest.raises(OSError, match="project payload disk is full"): + _ = evaluate_order_rules(project=project, root=tmp_path) + + assert time.monotonic() - failures[0] < test_case.expected_max_seconds + assert len(launches) == test_case.hosts + assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) + + +@pytest.mark.parametrize( + "test_case", + ( + AwaitInputsTestCase( + description="parent exited before publishing", + reported_parent_pid=1, + existing_files=(), + expected_error=RulesError, + expected_max_seconds=1.0, + ), + AwaitInputsTestCase( + description="run abandoned by its parent", + reported_parent_pid=4242, + existing_files=("project.abandoned",), + expected_error=HostCancelledError, + expected_max_seconds=1.0, + ), + ), + ids=lambda case: case.description, +) +def test_given_waiting_host_when_parent_exits_or_run_is_abandoned_then_wait_stops_promptly( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: AwaitInputsTestCase +) -> None: + touch_files(root=tmp_path, names=test_case.existing_files) + monkeypatch.setattr(custom_host.os, "getppid", lambda: test_case.reported_parent_pid) + started: float = time.monotonic() + + with pytest.raises(test_case.expected_error): + custom_host.await_inputs( + path=tmp_path / "project.pickle", + cancel_markers=(tmp_path / "project.abandoned",), + parent_pid=4242, + ) + + assert time.monotonic() - started < test_case.expected_max_seconds + + +@pytest.mark.parametrize( + "test_case", + ( + AwaitPublishedInputsTestCase( + description="payload published while the parent runs", + reported_parent_pid=4242, + existing_files=("project.pickle",), + expected_max_seconds=1.0, + ), + ), + ids=lambda case: case.description, +) +def test_given_published_payload_when_parent_still_runs_then_wait_returns( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: AwaitPublishedInputsTestCase +) -> None: + touch_files(root=tmp_path, names=test_case.existing_files) + monkeypatch.setattr(custom_host.os, "getppid", lambda: test_case.reported_parent_pid) + started: float = time.monotonic() + + custom_host.await_inputs( + path=tmp_path / "project.pickle", + cancel_markers=(tmp_path / "project.abandoned",), + parent_pid=4242, + ) + + assert time.monotonic() - started < test_case.expected_max_seconds + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/tests/unit/src/sqlbuild/rule_engine/main/evaluate_rule/_test_types.py b/tests/unit/src/sqlbuild/rule_engine/main/evaluate_rule/_test_types.py index 2b384ced2d..b5cd76213b 100644 --- a/tests/unit/src/sqlbuild/rule_engine/main/evaluate_rule/_test_types.py +++ b/tests/unit/src/sqlbuild/rule_engine/main/evaluate_rule/_test_types.py @@ -25,3 +25,11 @@ class EvaluateRuleDeterminismTestCase: description: str rule_case: RuleCase expected_error_pattern: str + + +@dataclass(frozen=True) +class AuthoringImportTestCase: + description: str + imported_module: str + checked_modules: tuple[str, ...] + expected_loaded_modules: tuple[str, ...] diff --git a/tests/unit/src/sqlbuild/rule_engine/main/evaluate_rule/test_authoring_imports.py b/tests/unit/src/sqlbuild/rule_engine/main/evaluate_rule/test_authoring_imports.py new file mode 100644 index 0000000000..7981434951 --- /dev/null +++ b/tests/unit/src/sqlbuild/rule_engine/main/evaluate_rule/test_authoring_imports.py @@ -0,0 +1,55 @@ +"""Importing the Rule authoring API stays cheap for every custom-rule host.""" + +import json +import subprocess +import sys + +import pytest + +from tests.unit.src.sqlbuild.rule_engine.main.evaluate_rule._test_types import ( + AuthoringImportTestCase, +) + + +@pytest.mark.parametrize( + "test_case", + ( + AuthoringImportTestCase( + description="rule modules import the authoring API", + imported_module="sqlbuild.rules", + checked_modules=( + "sqlbuild.rule_engine._helpers.engine.rule_harness", + "sqlbuild.compiler.pipeline.main.project", + "sqlbuild.adapters.duckdb.classes.duckdb_adapter", + ), + expected_loaded_modules=(), + ), + AuthoringImportTestCase( + description="hosts import the custom-rule host", + imported_module="sqlbuild.rule_engine._helpers.host.custom_host", + checked_modules=( + "sqlbuild.rule_engine._helpers.engine.rule_harness", + "sqlbuild.compiler.pipeline.main.project", + ), + expected_loaded_modules=(), + ), + ), + ids=lambda case: case.description, +) +def test_given_fresh_interpreter_when_importing_authoring_api_then_harness_is_not_loaded( + test_case: AuthoringImportTestCase, +) -> None: + script: str = ( + f"import json, sys\nimport {test_case.imported_module}\n" + f"print(json.dumps(sorted(set({list(test_case.checked_modules)!r}) & set(sys.modules))))" + ) + + result: subprocess.CompletedProcess[str] = subprocess.run( + [sys.executable, "-c", script], capture_output=True, text=True, check=True + ) + + assert tuple(json.loads(result.stdout)) == test_case.expected_loaded_modules + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-vv"])) From 358b33f88d233626b6ec47f224984f68cf82b8ef Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Tue, 6 Oct 2026 05:22:34 +0100 Subject: [PATCH 2/4] perf(rules): send custom-rule hosts only the project state Rule facts read The host payload drops model fields and project state that no Rule fact exposes, such as the second authored query copy, lineage columns, and the scope index, so the parent encodes and every host decodes less. Model fields are dropped while the payload is written, after hosts have started, so runs whose custom results are all cached pay nothing extra. --- .../_helpers/engine/custom_rules.py | 21 ++- .../_helpers/host/custom_host_pool.py | 27 +++- .../rule_engine/host_payload/_test_types.py | 8 ++ .../rule_engine/host_payload/helpers.py | 125 +++++++++++++++++ .../host_payload/test_host_payload_facts.py | 132 ++++++++++++++++++ 5 files changed, 304 insertions(+), 9 deletions(-) create mode 100644 tests/integration/src/sqlbuild/rule_engine/host_payload/_test_types.py create mode 100644 tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py create mode 100644 tests/integration/src/sqlbuild/rule_engine/host_payload/test_host_payload_facts.py diff --git a/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py b/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py index 9833dc1e18..1ec801614e 100644 --- a/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py +++ b/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py @@ -15,6 +15,7 @@ from sqlbuild.compiler.compile.models import CompiledProject from sqlbuild.compiler.fact_cache.main.code_identity import compiled_code_identity +from sqlbuild.compiler.scopes.models import ScopeIndex from sqlbuild.rule_engine._helpers.engine.custom_rule_evidence import ( custom_rule_implementation_fingerprint, custom_rule_import_closure, @@ -81,9 +82,10 @@ def evaluate_custom_rules_cached( custom: tuple[Rule, ...] = tuple(sorted(rules, key=lambda rule: rule.code)) if not custom: return CustomRulesOutcome(findings=(), cache_hits=0, cache_misses=0, custom_ms=0) + facts_project: CompiledProject = host_project(project) if not config.cache.enabled: uncached: CustomHostRun = run_custom_hosts( - project=_host_project(project), + project=facts_project, config=config, project_dir=project_dir, dialect=dialect, @@ -104,7 +106,7 @@ def evaluate_custom_rules_cached( custom_ms=_elapsed_ms(started), ) digests: FactDigests = FactDigests( - views=_views_factory(project=project, project_dir=project_dir, dialect=dialect) + views=_views_factory(project=facts_project, project_dir=project_dir, dialect=dialect) ) subjects: dict[str, tuple[_Subject, ...]] = _subjects(project=project, rules=custom) identities: dict[str, str] = { @@ -150,7 +152,7 @@ def evaluate_custom_rules_cached( ) if plan: run: CustomHostRun = run_custom_hosts( - project=_host_project(project), + project=facts_project, config=config, project_dir=project_dir, dialect=dialect, @@ -187,9 +189,7 @@ def _views_factory( *, project: CompiledProject, project_dir: Path, dialect: str ) -> Callable[[], RuleFactViews]: def build() -> RuleFactViews: - return build_rule_fact_views( - project=_host_project(project), project_dir=project_dir, dialect=dialect - ) + return build_rule_fact_views(project=project, project_dir=project_dir, dialect=dialect) return build @@ -268,7 +268,9 @@ def _readset_digest( return memo[index] -def _host_project(project: CompiledProject) -> CompiledProject: +def host_project(project: CompiledProject) -> CompiledProject: + """Drop project state no Rule fact reads, before digests and host payloads use it.""" + return replace( project, binding_catalog=None, @@ -278,6 +280,11 @@ def _host_project(project: CompiledProject) -> CompiledProject: hook_functions=(), materialization_files=(), external_sql_reference_resolver=None, + sql_scenarios=(), + sql_hook_files=(), + diagnostics=(), + unmatched_literal_sql_relations=(), + scope_index=ScopeIndex(), ) diff --git a/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py b/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py index 8d5710be3a..cff549221a 100644 --- a/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py +++ b/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py @@ -17,7 +17,7 @@ import orjson import sqlbuild._native as _native -from sqlbuild.compiler.compile.models import CompiledProject +from sqlbuild.compiler.compile.models import CompiledModel, CompiledProject from sqlbuild.rule_engine._helpers.host.host_capacity import available_cores from sqlbuild.rule_engine.constants import ( CUSTOM_HOST_ABANDONED_SUFFIX, @@ -99,12 +99,13 @@ def run_custom_hosts( def _write_inputs(*, path: Path, project: CompiledProject, config: RulesConfig) -> None: """Publish the project payload atomically; hosts started earlier wait for it to appear.""" + payload: CompiledProject = host_payload_project(project) with tempfile.NamedTemporaryFile( dir=path.parent, prefix="writing-", suffix=".pickle", delete=False ) as handle: partial: Path = Path(handle.name) try: - pickle.dump((project, config), handle) + pickle.dump((payload, config), handle) except BaseException: handle.close() partial.unlink(missing_ok=True) @@ -112,6 +113,28 @@ def _write_inputs(*, path: Path, project: CompiledProject, config: RulesConfig) partial.replace(path) +def host_payload_project(project: CompiledProject) -> CompiledProject: + """Return the project hosts decode, without model state no Rule fact reads.""" + + return replace(project, models=tuple(map(_payload_model, project.models))) + + +def _payload_model(model: CompiledModel) -> CompiledModel: + """Drop model state no Rule fact reads, so hosts decode less of every model.""" + + return replace( + model, + references=(), + fast_lineage_columns=None, + authored_query_sql="", + output_column_locations={}, + macro_deps=(), + enum_columns={}, + binding_diagnostics=(), + unchecked_output_columns=frozenset(), + ) + + def _abandon_inputs(marker: Path) -> bool: """Tell waiting hosts no payload will come; report whether they can see the signal.""" diff --git a/tests/integration/src/sqlbuild/rule_engine/host_payload/_test_types.py b/tests/integration/src/sqlbuild/rule_engine/host_payload/_test_types.py new file mode 100644 index 0000000000..63224ceda7 --- /dev/null +++ b/tests/integration/src/sqlbuild/rule_engine/host_payload/_test_types.py @@ -0,0 +1,8 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class HostPayloadFactsTestCase: + description: str + project_files: dict[str, str] + expected_unpopulated_trimmed_fields: frozenset[str] diff --git a/tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py b/tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py new file mode 100644 index 0000000000..aef813813c --- /dev/null +++ b/tests/integration/src/sqlbuild/rule_engine/host_payload/helpers.py @@ -0,0 +1,125 @@ +import dataclasses +from itertools import product +from pathlib import Path +from typing import Any + +import pytest + +from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter +from sqlbuild.compiler.compile.models import CompiledModel, CompiledProject +from sqlbuild.compiler.discovery.main.discover import discover_project_inputs +from sqlbuild.compiler.pipeline.main.project import compile_project +from sqlbuild.rule_engine._helpers.engine import custom_rules, fact_replay +from sqlbuild.rule_engine._helpers.engine.fact_replay import ( + fact_outcome_digest, + shared_fact_encoder, +) +from sqlbuild.rule_engine._helpers.host import custom_host_pool +from sqlbuild.rule_engine.classes.project_tree import public_model +from sqlbuild.rule_engine.classes.rule_context import RuleFactViews, build_rule_fact_views +from sqlbuild.rule_engine.models import Model +from sqlbuild.rule_engine.types import FactKey + +_TREE_ARGUMENTS: tuple[str, ...] = ("models/marts/order_totals.sql", "models") + + +class OrdersReferenceResolver: + """External reference resolver that knows no extra relations.""" + + def validate_model_names(self, *, known_model_names: set[str]) -> None: + del known_model_names + + def extend_sql_test_model_names(self, *, known_model_names: set[str]) -> set[str]: + return known_model_names + + def extend_sql_test_source_names(self, *, known_source_names: set[str]) -> set[str]: + return known_source_names + + def extend_sql_test_seed_names(self, *, known_seed_names: set[str]) -> set[str]: + return known_seed_names + + def validate_reference( + self, + *, + ref_kind: str, + ref_name: str, + ref_package: str | None, + owner_relative_sql_path: Path, + ) -> None: + del ref_kind, ref_name, ref_package, owner_relative_sql_path + + def resolve_reference( + self, *, ref_kind: str, ref_name: str, ref_package: str | None + ) -> str | None: + del ref_kind, ref_name, ref_package + return None + + +def compile_files(*, root: Path, files: dict[str, str]) -> CompiledProject: + """Write and compile one DuckDB project with an external reference resolver.""" + + for relative_path, contents in files.items(): + path: Path = root / relative_path + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(contents, encoding="utf-8") + return compile_project( + discovered_inputs=discover_project_inputs(project_dir=root), + adapter=DuckDbAdapter(), + external_sql_reference_resolver=OrdersReferenceResolver(), + ) + + +def dropped_fields(*, before: object, after: object) -> frozenset[str]: + """Return dataclass fields whose values differ between two snapshots of one object.""" + + names: tuple[str, ...] = tuple(field.name for field in dataclasses.fields(type(before))) + return frozenset(filter(lambda name: getattr(before, name) != getattr(after, name), names)) + + +def every_fact_key(project: CompiledProject) -> tuple[FactKey, ...]: + """Return a key for every fact Rules can read, over every model of the project.""" + + models: tuple[Model, ...] = tuple(public_model(model) for model in project.models) + return ( + *((fact,) for fact in fact_replay._GLOBAL_FACTS), + *( + (fact, model.path.as_posix()) + for fact, model in product(fact_replay._MODEL_PATH_FACTS, models) + ), + *((fact, model.name) for fact, model in product(fact_replay._MODEL_NAME_FACTS, models)), + *( + (fact, *_TREE_ARGUMENTS[: fact_replay._TEXT_FACT_ARITY[fact]]) + for fact in fact_replay._TEXT_FACTS + ), + ) + + +def fact_digests(*, project: CompiledProject, root: Path) -> dict[FactKey, tuple[str, str]]: + """Digest every fact as hosts observe it and as whole-project digests encode it.""" + + views: RuleFactViews = build_rule_fact_views( + project=project, project_dir=root, dialect="duckdb" + ) + return { + key: ( + fact_outcome_digest(views=views, key=key), + fact_outcome_digest(views=views, key=key, shared=shared_fact_encoder()), + ) + for key in every_fact_key(project) + } + + +def record_trimmed_fields( + monkeypatch: pytest.MonkeyPatch, +) -> dict[type[object], set[str]]: + """Record which fields the host-payload trimmers replace, by dataclass type.""" + + trimmed: dict[type[object], set[str]] = {CompiledProject: set(), CompiledModel: set()} + + def recording_replace(obj: Any, /, **changes: object) -> Any: + trimmed.setdefault(type(obj), set()).update(changes) + return dataclasses.replace(obj, **changes) + + monkeypatch.setattr(custom_rules, "replace", recording_replace) + monkeypatch.setattr(custom_host_pool, "replace", recording_replace) + return trimmed diff --git a/tests/integration/src/sqlbuild/rule_engine/host_payload/test_host_payload_facts.py b/tests/integration/src/sqlbuild/rule_engine/host_payload/test_host_payload_facts.py new file mode 100644 index 0000000000..8b65e73896 --- /dev/null +++ b/tests/integration/src/sqlbuild/rule_engine/host_payload/test_host_payload_facts.py @@ -0,0 +1,132 @@ +"""Every Rule fact reads the same values from the trimmed host payload as from the full project.""" + +from pathlib import Path + +import pytest + +from sqlbuild.compiler.compile.models import CompiledModel, CompiledProject +from sqlbuild.rule_engine._helpers.engine.custom_rules import host_project +from sqlbuild.rule_engine._helpers.host.custom_host_pool import host_payload_project +from sqlbuild.rule_engine.classes.project_tree import public_model +from sqlbuild.rule_engine.types import FactKey +from tests.integration.src.sqlbuild.rule_engine.host_payload._test_types import ( + HostPayloadFactsTestCase, +) +from tests.integration.src.sqlbuild.rule_engine.host_payload.helpers import ( + compile_files, + dropped_fields, + fact_digests, + record_trimmed_fields, +) + + +@pytest.mark.parametrize( + "test_case", + ( + HostPayloadFactsTestCase( + description="sources, refs, contracts, macros, declarations, tests, and audits", + project_files={ + "sqlbuild_project.toml": 'name = "orders"\nadapter = "duckdb"\n', + "sources/raw.yml": ( + "sources:\n" + " - name: raw_orders\n description: Raw orders.\n" + ' expression: "(SELECT 1 AS id, CAST(10.5 AS DOUBLE) AS amount)"\n' + " columns:\n - name: id\n type: INTEGER\n" + " - name: raw_customers\n description: Raw customers.\n" + ' expression: "(SELECT 1 AS customer_id)"\n' + ), + "constants/policy.sql": "CONSTANT (name minimum_amount, value 7);\n", + "enums/order_status.sql": "ENUM (name order_status, members [PAID, OPEN]);\n", + "models/marts/_sqlbuild/_macros/amounts.py": ( + 'def doubled(amount: str) -> str:\n return f"({amount}) * 2"\n' + ), + "models/staging/stg_orders.sql": ( + "MODEL (description 'Staged orders.',\n contract enforced,\n" + " columns (\n" + " order_id (type INTEGER, nullable false, audits [not_null]),\n" + " amount (type DOUBLE),\n status (type order_status),\n ),\n);\n\n" + "SELECT CAST(id AS INTEGER) AS order_id, amount, 'PAID' AS status\n" + 'FROM __source("raw_orders")\n' + ), + "models/marts/order_totals.sql": ( + 'MODEL (description "Order totals", materialized table);\n\n' + 'SELECT o.order_id, @doubled("o.amount") AS doubled_amount, o.status\n' + 'FROM __ref("stg_orders") AS o WHERE o.amount > 7\n' + ), + "models/marts/order_flags.sql": ( + 'MODEL (description "Order flags");\n\n' + 'SELECT o.order_id, o.missing_flag AS flag FROM __ref("stg_orders") AS o\n' + ), + "models/marts/flag_counts.sql": ( + 'MODEL (description "Flag counts");\n\n' + 'SELECT f.flag FROM __ref("order_flags") AS f\n' + ), + "hooks/sql/analyze_orders.sql": ( + 'HOOK (description "Analyze orders");\n\nSELECT 1\n' + ), + "hooks/python/order_hooks.py": ( + "from sqlbuild.hooks import hook\n\n" + "@hook\n" + "def log_orders(ctx):\n" + " '''Log orders.'''\n" + " ctx.execute_sql('INSERT INTO audit_log SELECT 1')\n" + ), + "python/loaders/order_rows.py": ( + "from sqlbuild.loaders import loader\n\n" + "@loader(columns=[{'name': 'order_id', 'type': 'INTEGER'}])\n" + "def fetch_orders(ctx):\n" + " '''Fetch orders.'''\n" + " return [{'order_id': 1}]\n" + ), + "materializations/copy_table.py": ( + "from sqlbuild.executor.custom.models import (\n" + " MaterializationContext,\n MaterializationResult,\n)\n\n\n" + "def materialize(ctx: MaterializationContext) -> MaterializationResult:\n" + ' ctx.execute_sql(f"CREATE TABLE {ctx.destination} AS {ctx.sql}")\n' + " return MaterializationResult(relation=ctx.destination)\n" + ), + "tests/scenarios/paid_orders.sql": ( + "SCENARIO (description 'Paid orders.');\n\n" + "WITH\n__source__raw_orders AS (SELECT 1 AS id, 10.5 AS amount),\n" + "__expected__stg_orders AS (\n" + " SELECT 1 AS order_id, 10.5 AS amount, 'PAID' AS status\n)\n" + "SELECT 1\n" + ), + "tests/unit/test_order_totals.sql": ( + "TEST ();\n\nWITH\n__ref__stg_orders AS (\n" + " SELECT 1 AS order_id, 10.0 AS amount\n),\n" + "__expected__order_totals AS (\n" + " SELECT 1 AS order_id, 20.0 AS doubled_amount, 'PAID' AS status\n)\n" + "SELECT 1\n" + ), + }, + expected_unpopulated_trimmed_fields=frozenset(), + ), + ), + ids=lambda case: case.description, +) +def test_given_compiled_project_when_trimming_host_payload_then_every_fact_is_unchanged( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostPayloadFactsTestCase +) -> None: + project: CompiledProject = compile_files(root=tmp_path, files=test_case.project_files) + trimmed_fields: dict[type[object], set[str]] = record_trimmed_fields(monkeypatch) + trimmed: CompiledProject = host_payload_project(host_project(project)) + + full_digests: dict[FactKey, tuple[str, str]] = fact_digests(project=project, root=tmp_path) + trimmed_digests: dict[FactKey, tuple[str, str]] = fact_digests(project=trimmed, root=tmp_path) + + assert trimmed_digests == full_digests + assert tuple(map(public_model, trimmed.models)) == tuple(map(public_model, project.models)) + assert ( + dropped_fields(before=project, after=trimmed) + == trimmed_fields[CompiledProject] - test_case.expected_unpopulated_trimmed_fields + ) + assert ( + frozenset().union( + *( + dropped_fields(before=model, after=trimmed_model) + for model, trimmed_model in zip(project.models, trimmed.models, strict=True) + ) + ) + == trimmed_fields[CompiledModel] - test_case.expected_unpopulated_trimmed_fields + ) From fe5d292276ecc45368e3e42bfd9247726f8b3a75 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Tue, 6 Oct 2026 05:22:34 +0100 Subject: [PATCH 3/4] perf(rules): resolve project-tree reads through cached directories Reading a regular, non-symlink file resolves its directory once per tree and checks the file itself with one lstat. Symlinks, missing files, and directories still take the full resolution, so outcomes and errors match. --- .../rule_engine/classes/project_tree.py | 33 +++++++++- .../classes/project_tree/_test_types.py | 9 +++ .../classes/project_tree/helpers.py | 30 +++++++++ .../classes/project_tree/test_read_text.py | 66 +++++++++++++++++++ 4 files changed, 135 insertions(+), 3 deletions(-) create mode 100644 tests/unit/src/sqlbuild/rule_engine/classes/project_tree/_test_types.py create mode 100644 tests/unit/src/sqlbuild/rule_engine/classes/project_tree/helpers.py create mode 100644 tests/unit/src/sqlbuild/rule_engine/classes/project_tree/test_read_text.py diff --git a/src/sqlbuild/rule_engine/classes/project_tree.py b/src/sqlbuild/rule_engine/classes/project_tree.py index 5bec6e4632..88f38ae34f 100644 --- a/src/sqlbuild/rule_engine/classes/project_tree.py +++ b/src/sqlbuild/rule_engine/classes/project_tree.py @@ -4,6 +4,7 @@ import fnmatch import os +import stat from pathlib import Path from sqlbuild.compiler.compile.models import CompiledProject @@ -21,6 +22,7 @@ class ProjectTree: def __init__(self, *, project_dir: Path, project: CompiledProject) -> None: self._root: Path = project_dir.resolve() self._discovered: tuple[ProjectPath, ...] | None = None + self._directories: dict[Path, Path | None] = {} self._resource_paths: dict[str, Model] = { item.relative_path.as_posix(): public_model(item) for item in project.models } @@ -70,21 +72,46 @@ def resources_under(self, path: str) -> tuple[Model, ...]: def read_text(self, path: str) -> str: self._validate_relative(path) - candidate: Path = (self._root / path).resolve() - if not candidate.is_relative_to(self._root) or not candidate.is_file(): + regular: Path | None = self._regular_file(path) + candidate: Path = (self._root / path).resolve() if regular is None else regular + if regular is None and ( + not candidate.is_relative_to(self._root) or not candidate.is_file() + ): raise RuleUsageError(f"project file does not exist: {path}") if candidate.suffix.lower() not in _TRACKED_TEXT_SUFFIXES: supported: str = ", ".join(sorted(_TRACKED_TEXT_SUFFIXES)) raise RuleUsageError( f"compiler rules input must be a tracked file type ({supported}): {path}" ) - if candidate.is_symlink(): + if regular is None and candidate.is_symlink(): raise RuleUsageError(f"project file must not be a symlink: {path}") try: return candidate.read_text(encoding="utf-8") except (OSError, UnicodeError) as error: raise RuleUsageError(f"could not read project file {path}: {error}") from error + def _regular_file(self, path: str) -> Path | None: + """Resolve a regular, non-symlink file under the project through its cached directory.""" + + relative: Path = Path(path) + if not relative.name: + return None + directory: Path | None + if relative.parent in self._directories: + directory = self._directories[relative.parent] + else: + resolved: Path = (self._root / relative.parent).resolve() + directory = resolved if resolved.is_relative_to(self._root) else None + self._directories[relative.parent] = directory + if directory is None: + return None + candidate: Path = directory / relative.name + try: + regular: bool = stat.S_ISREG(os.lstat(candidate).st_mode) + except OSError: + return None + return candidate if regular else None + def _discover(self) -> tuple[ProjectPath, ...]: nodes: list[ProjectPath] = [] for root, directories, filenames in os.walk(self._root, followlinks=False): diff --git a/tests/unit/src/sqlbuild/rule_engine/classes/project_tree/_test_types.py b/tests/unit/src/sqlbuild/rule_engine/classes/project_tree/_test_types.py new file mode 100644 index 0000000000..99914f928d --- /dev/null +++ b/tests/unit/src/sqlbuild/rule_engine/classes/project_tree/_test_types.py @@ -0,0 +1,9 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class ReadTextTestCase: + description: str + path: str + expected_text: str | None = None + expected_error: str = "" diff --git a/tests/unit/src/sqlbuild/rule_engine/classes/project_tree/helpers.py b/tests/unit/src/sqlbuild/rule_engine/classes/project_tree/helpers.py new file mode 100644 index 0000000000..387bc9c79d --- /dev/null +++ b/tests/unit/src/sqlbuild/rule_engine/classes/project_tree/helpers.py @@ -0,0 +1,30 @@ +from pathlib import Path + +from sqlbuild.compiler.compile.models import CompiledProject +from sqlbuild.rule_engine.classes.project_tree import ProjectTree +from tests.unit.src.sqlbuild.rule_engine.main.evaluate.helpers import build_project + +ORDERS_SQL: str = "SELECT 1 AS order_id\n" +LIMITS_YAML: str = "max_orders: 3\n" + + +def linked_project_tree(tmp_path: Path) -> ProjectTree: + """Return a project tree with regular files and symlinks inside and outside the project.""" + + root: Path = tmp_path / "orders" + outside: Path = tmp_path / "outside" + (root / "models").mkdir(parents=True) + (root / "rules" / "inputs").mkdir(parents=True) + outside.mkdir() + (root / "models" / "orders.sql").write_text(ORDERS_SQL, encoding="utf-8") + (root / "models" / "notes.txt").write_text("notes\n", encoding="utf-8") + (root / "rules" / "inputs" / "limits.yaml").write_text(LIMITS_YAML, encoding="utf-8") + (outside / "customers.sql").write_text("SELECT 1 AS customer_id\n", encoding="utf-8") + (root / "models" / "orders_alias.sql").symlink_to(root / "models" / "orders.sql") + (root / "models" / "customers.sql").symlink_to(outside / "customers.sql") + (root / "linked_models").symlink_to(root / "models", target_is_directory=True) + (root / "outside_models").symlink_to(outside, target_is_directory=True) + project: CompiledProject = build_project( + name="orders", relative_path="models/orders.sql", sql=ORDERS_SQL, config_values={} + ) + return ProjectTree(project_dir=root, project=project) diff --git a/tests/unit/src/sqlbuild/rule_engine/classes/project_tree/test_read_text.py b/tests/unit/src/sqlbuild/rule_engine/classes/project_tree/test_read_text.py new file mode 100644 index 0000000000..addaa83265 --- /dev/null +++ b/tests/unit/src/sqlbuild/rule_engine/classes/project_tree/test_read_text.py @@ -0,0 +1,66 @@ +"""Project-tree reads resolve every path exactly as a fully resolved read would.""" + +from pathlib import Path + +import pytest + +from sqlbuild.rule_engine.classes.project_tree import ProjectTree +from sqlbuild.rule_engine.exceptions import RuleUsageError +from tests.unit.src.sqlbuild.rule_engine.classes.project_tree._test_types import ReadTextTestCase +from tests.unit.src.sqlbuild.rule_engine.classes.project_tree.helpers import ( + LIMITS_YAML, + ORDERS_SQL, + linked_project_tree, +) + + +@pytest.mark.parametrize( + "test_case", + ( + ReadTextTestCase("regular file", "models/orders.sql", expected_text=ORDERS_SQL), + ReadTextTestCase("nested file", "rules/inputs/limits.yaml", expected_text=LIMITS_YAML), + ReadTextTestCase("dot prefix", "./models/orders.sql", expected_text=ORDERS_SQL), + ReadTextTestCase("symlinked file", "models/orders_alias.sql", expected_text=ORDERS_SQL), + ReadTextTestCase("symlinked folder", "linked_models/orders.sql", expected_text=ORDERS_SQL), + ), + ids=lambda case: case.description, +) +def test_given_readable_project_file_when_reading_text_twice_then_returns_its_contents( + tmp_path: Path, test_case: ReadTextTestCase +) -> None: + tree: ProjectTree = linked_project_tree(tmp_path) + + first: str = tree.read_text(test_case.path) + second: str = tree.read_text(test_case.path) + + assert first == second == test_case.expected_text + + +@pytest.mark.parametrize( + "test_case", + ( + ReadTextTestCase("missing file", "models/missing.sql", expected_error="does not exist"), + ReadTextTestCase("directory", "models", expected_error="does not exist"), + ReadTextTestCase("project root", "", expected_error="does not exist"), + ReadTextTestCase("beneath a file", "models/orders.sql/x.sql", expected_error="not exist"), + ReadTextTestCase( + "folder linked outside", "outside_models/customers.sql", expected_error="not exist" + ), + ReadTextTestCase("file linked outside", "models/customers.sql", expected_error="not exist"), + ReadTextTestCase("untracked suffix", "models/notes.txt", expected_error="tracked file"), + ReadTextTestCase("parent escape", "../outside/customers.sql", expected_error="escapes"), + ), + ids=lambda case: case.description, +) +def test_given_unreadable_project_path_when_reading_text_then_raises_usage_error( + tmp_path: Path, test_case: ReadTextTestCase +) -> None: + tree: ProjectTree = linked_project_tree(tmp_path) + _ = tree.read_text("models/orders.sql") + + with pytest.raises(RuleUsageError, match=test_case.expected_error): + _ = tree.read_text(test_case.path) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-vv"])) From 3d69717ea8bdc1f69d6eb7f1603e357bedebc32d Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Tue, 6 Oct 2026 05:22:34 +0100 Subject: [PATCH 4/4] perf(rules): digest shared whole-project facts once Whole-project facts such as project.sources repeat objects shared between resources: every source declared in one file carries that file's contents and entries. Digests of these facts now encode each compiler dataclass once and refer to it by content digest, so the digest grows with distinct content instead of with repetitions. Model and project-tree facts keep the digest a host observes. A real-CLI edit chain with model- and project-subject custom rules split across hosts compares every warm compile with a cache-free oracle, and the oracle is shown to fail when source digests are deliberately stale. --- .../_helpers/engine/fact_replay.py | 19 ++- .../rule_engine/classes/fact_digests.py | 9 +- .../classes/shared_fact_encoder.py | 41 ++++++ src/sqlbuild/rule_engine/constants.py | 1 + .../cli/commands/main/rules/_test_types.py | 10 ++ .../cli/commands/main/rules/helpers.py | 121 +++++++++++++++- .../test_custom_rules_incremental_sequence.py | 136 ++++++++++++++++++ .../shared_fact_encoder/_test_types.py | 16 +++ .../classes/shared_fact_encoder/helpers.py | 60 ++++++++ .../test_shared_fact_encoder.py | 108 ++++++++++++++ 10 files changed, 516 insertions(+), 5 deletions(-) create mode 100644 src/sqlbuild/rule_engine/classes/shared_fact_encoder.py create mode 100644 tests/e2e/src/sqlbuild/cli/commands/main/rules/test_custom_rules_incremental_sequence.py create mode 100644 tests/unit/src/sqlbuild/rule_engine/classes/shared_fact_encoder/_test_types.py create mode 100644 tests/unit/src/sqlbuild/rule_engine/classes/shared_fact_encoder/helpers.py create mode 100644 tests/unit/src/sqlbuild/rule_engine/classes/shared_fact_encoder/test_shared_fact_encoder.py diff --git a/src/sqlbuild/rule_engine/_helpers/engine/fact_replay.py b/src/sqlbuild/rule_engine/_helpers/engine/fact_replay.py index 7404be19f1..0832f7e1a2 100644 --- a/src/sqlbuild/rule_engine/_helpers/engine/fact_replay.py +++ b/src/sqlbuild/rule_engine/_helpers/engine/fact_replay.py @@ -9,6 +9,7 @@ import orjson +from sqlbuild.rule_engine.classes.shared_fact_encoder import SharedFactEncoder from sqlbuild.rule_engine.classes.sql_document import SqlDocument from sqlbuild.rule_engine.constants import ( FACT_AUDITS_ALL, @@ -98,6 +99,7 @@ FACT_TREE_READ_TEXT: 1, } _TREE_FACTS: frozenset[str] = frozenset({FACT_TREE_PATHS, *_TEXT_FACTS}) +_SHARED_FACTS: frozenset[str] = frozenset(_GLOBAL_FACTS) - _TREE_FACTS def fact_key_payload(key: tuple[object, ...]) -> list[str]: @@ -123,8 +125,10 @@ def full_fact_keys(views: RuleFactViews) -> tuple[FactKey, ...]: return tuple(keys) -def fact_outcome_digest(*, views: RuleFactViews, key: FactKey) -> str: - """Recompute one recorded fact and digest its value or raised error.""" +def fact_outcome_digest( + *, views: RuleFactViews, key: FactKey, shared: SharedFactEncoder | None = None +) -> str: + """Recompute one recorded fact and digest its value, sharing whole-project encodings.""" try: value: object = _fact_value(views=views, key=key) @@ -132,9 +136,20 @@ def fact_outcome_digest(*, views: RuleFactViews, key: FactKey) -> str: raise except Exception as error: return fact_error_digest(error) + if shared is not None and key[0] in _SHARED_FACTS: + try: + return hashlib.sha256(shared.encode(("value", value))).hexdigest() + except (TypeError, FactDigestError) as error: + raise FactDigestError(str(error)) from error return fact_value_digest(value) +def shared_fact_encoder() -> SharedFactEncoder: + """Return an encoder for whole-project facts that share compiler objects.""" + + return SharedFactEncoder(options=_ORJSON_OPTIONS, encode_default=_encode_default) + + def fact_value_digest(value: object) -> str: """Digest one fact value exactly as a recomputed fact would digest it.""" diff --git a/src/sqlbuild/rule_engine/classes/fact_digests.py b/src/sqlbuild/rule_engine/classes/fact_digests.py index 05c663cc44..521eae9dd6 100644 --- a/src/sqlbuild/rule_engine/classes/fact_digests.py +++ b/src/sqlbuild/rule_engine/classes/fact_digests.py @@ -8,7 +8,11 @@ import orjson -from sqlbuild.rule_engine._helpers.engine.fact_replay import fact_outcome_digest +from sqlbuild.rule_engine._helpers.engine.fact_replay import ( + fact_outcome_digest, + shared_fact_encoder, +) +from sqlbuild.rule_engine.classes.shared_fact_encoder import SharedFactEncoder from sqlbuild.rule_engine.models import RuleFactViews from sqlbuild.rule_engine.types import FactKey @@ -20,6 +24,7 @@ def __init__(self, *, views: Callable[[], RuleFactViews]) -> None: self._views_factory: Callable[[], RuleFactViews] = views self._views: RuleFactViews | None = None self._digests: dict[FactKey, str] = {} + self._shared: SharedFactEncoder = shared_fact_encoder() @property def views(self) -> RuleFactViews: @@ -33,7 +38,7 @@ def digest(self, key: FactKey) -> str: cached: str | None = self._digests.get(key) if cached is not None: return cached - digest: str = fact_outcome_digest(views=self.views, key=key) + digest: str = fact_outcome_digest(views=self.views, key=key, shared=self._shared) self._digests[key] = digest return digest diff --git a/src/sqlbuild/rule_engine/classes/shared_fact_encoder.py b/src/sqlbuild/rule_engine/classes/shared_fact_encoder.py new file mode 100644 index 0000000000..6ce301a3bf --- /dev/null +++ b/src/sqlbuild/rule_engine/classes/shared_fact_encoder.py @@ -0,0 +1,41 @@ +"""Digest whole-project fact values whose resources share large compiler objects.""" + +from __future__ import annotations + +import dataclasses +import hashlib +from collections.abc import Callable + +import orjson + +from sqlbuild.rule_engine.constants import SHARED_FACT_DATACLASS_TAG + + +class SharedFactEncoder: + """Encode each compiler dataclass once and refer to it by digest wherever it repeats.""" + + def __init__(self, *, options: int, encode_default: Callable[[object], object]) -> None: + self._options: int = options | orjson.OPT_PASSTHROUGH_DATACLASS + self._encode_default: Callable[[object], object] = encode_default + self._fields: dict[type, tuple[str, ...]] = {} + self._digests: dict[int, tuple[object, str]] = {} + + def encode(self, value: object) -> bytes: + return orjson.dumps(value, default=self._default, option=self._options) + + def _default(self, value: object) -> object: + kind: type = type(value) + names: tuple[str, ...] | None = self._fields.get(kind) + if names is None: + if not dataclasses.is_dataclass(kind): + return self._encode_default(value) + names = tuple(field.name for field in dataclasses.fields(kind)) + self._fields[kind] = names + cached: tuple[object, str] | None = self._digests.get(id(value)) + if cached is None: + encoded: bytes = self.encode( + [f"{kind.__module__}.{kind.__qualname__}", [getattr(value, name) for name in names]] + ) + cached = (value, hashlib.sha256(encoded).hexdigest()) + self._digests[id(value)] = cached + return [SHARED_FACT_DATACLASS_TAG, cached[1]] diff --git a/src/sqlbuild/rule_engine/constants.py b/src/sqlbuild/rule_engine/constants.py index d8ba681b11..c3fe200ed5 100644 --- a/src/sqlbuild/rule_engine/constants.py +++ b/src/sqlbuild/rule_engine/constants.py @@ -16,6 +16,7 @@ NATIVE_RULES_MEMO_VERSION: str = "native-rules-response-v2" NATIVE_RULES_MEMO_FILE: str = "target/rules-cache/bulk/native-response.json" CUSTOM_HOST_INPUT_TUPLE_SIZE: int = 2 +SHARED_FACT_DATACLASS_TAG: str = "\0dataclass" CUSTOM_HOST_ABANDONED_SUFFIX: str = ".abandoned" CUSTOM_HOST_INPUT_POLL_SECONDS: float = 0.002 CUSTOM_HOST_INPUT_WAIT_SECONDS: float = 120.0 diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/rules/_test_types.py b/tests/e2e/src/sqlbuild/cli/commands/main/rules/_test_types.py index 0c2866e28f..4ca693ff57 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/rules/_test_types.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/rules/_test_types.py @@ -222,3 +222,13 @@ class RulesEditChainCase: edits: tuple[RulesCacheEditCase, ...] expected_cold_codes: tuple[str, ...] expected_exit_code: int = 1 + + +@dataclass(frozen=True) +class BrokenInvalidationChainCase: + """An edit chain compiled with deliberately stale fact digests against a cache-free oracle.""" + + description: str + edits: tuple[RulesCacheEditCase, ...] + stale_digests_sitecustomize: str + expected_first_divergent_edit: str diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/rules/helpers.py b/tests/e2e/src/sqlbuild/cli/commands/main/rules/helpers.py index c6b8081c7a..82f8007dd0 100644 --- a/tests/e2e/src/sqlbuild/cli/commands/main/rules/helpers.py +++ b/tests/e2e/src/sqlbuild/cli/commands/main/rules/helpers.py @@ -3,6 +3,7 @@ import hashlib import json import os +import shutil import subprocess import sys from collections.abc import Callable @@ -411,6 +412,121 @@ def flagged_values(*, model: Model, ctx: RuleContext) -> list[Finding]: } +_ORDER_SLICE_COUNT: int = 45 +CUSTOM_RULES_EDIT_PROJECT: dict[str, str] = { + "sqlbuild_project.toml": ( + 'name = "orders"\nadapter = "duckdb"\n\n' + '[connection]\ndatabase = "warehouse.duckdb"\n\n' + '[rules]\nselect = ["XSQBR"]\n\n' + "[rules.thresholds]\nmin_custom_rule_test_cases = 0\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, 1 AS customer_id)"\n' + " - name: raw_customers\n description: Raw customers.\n" + ' expression: "(SELECT 1 AS customer_id)"\n' + " - name: raw_products\n description: Raw products.\n" + ' expression: "(SELECT 1 AS product_id)"\n' + ), + "models/staging/stg_orders.sql": ( + 'MODEL (description "Staged orders");\n\n' + 'SELECT id AS order_id, amount, customer_id FROM __source("raw_orders")\n' + ), + "models/staging/stg_customers.sql": ( + 'MODEL (description "Staged customers");\n\n' + 'SELECT customer_id FROM __source("raw_customers")\n' + ), + "models/marts/order_totals.sql": ( + 'MODEL (description "Order totals");\n\n' + 'SELECT order_id, amount * 2 AS doubled_amount FROM __ref("stg_orders") WHERE amount > 7\n' + ), + **{ + f"models/intermediate/orders_{index:03d}.sql": ( + f'MODEL (description "Order slice {index}");\n\n' + 'SELECT order_id FROM __ref("stg_orders")\n' + ) + for index in range(_ORDER_SLICE_COUNT) + }, + "rules/approved_sources.yaml": "raw_orders\nraw_customers\n", + "rules/orders.py": """from sqlbuild.rules import Finding, Model, Project, RuleContext, rule + + +@rule(code="XSQBRORD101", message="Filters hard-code an amount", remediation="Use a constant.") +def high_amount_filters(*, model: Model, ctx: RuleContext) -> list[Finding]: + expanded: str = ctx.sql.for_model(model).expanded.source + return [ctx.finding(subject=model)] if "amount > 9" in expanded else [] + + +@rule(code="XSQBRORD102", message="Staging models need a consumer", remediation="Use it.") +def staging_consumers(*, model: Model, ctx: RuleContext) -> list[Finding]: + if model.path.parent.name != "staging": + return [] + return [] if ctx.graph.dependents(model) else [ctx.finding(subject=model)] + + +@rule(code="XSQBRORD103", message="Marts must enforce a contract", remediation="Enforce it.") +def mart_contracts(*, model: Model, ctx: RuleContext) -> list[Finding]: + if model.path.parent.name != "marts": + return [] + return [] if ctx.contracts.enforced(model) else [ctx.finding(subject=model)] + + +@rule(code="XSQBRORD201", message="Sources must be approved", remediation="Approve it.") +def approved_sources(*, project: Project, ctx: RuleContext) -> list[Finding]: + del project + approved: set[str] = set(ctx.project.tree.read_text("rules/approved_sources.yaml").split()) + names: list[str] = sorted(str(getattr(source, "name")) for source in ctx.project.sources) + return [ + ctx.finding(subject="sources/raw.yml", message=f"Source {name} is not approved") + for name in names + if name not in approved + ] + + +@rule(code="XSQBRORD202", message="Marts must stay few", remediation="Consolidate them.") +def few_marts(*, project: Project, ctx: RuleContext) -> list[Finding]: + del project + marts: list[Model] = [item for item in ctx.project.models if item.path.parent.name == "marts"] + if len(marts) < 2: + return [] + return [ctx.finding(subject=marts[0], message=f"{len(marts)} marts")] +""", +} +STALE_SOURCE_FACTS_SITECUSTOMIZE: str = """from sqlbuild.rule_engine.classes import fact_digests + +_digest = fact_digests.FactDigests.digest + + +def _stale_sources(self, key): + return "stale" if key == ("project.sources",) else _digest(self, key) + + +fact_digests.FactDigests.digest = _stale_sources +""" + + +def custom_rules_edit_chain( + *, root: Path, edits: tuple[Any, ...], environment: tuple[tuple[str, str], ...] = () +) -> tuple[ + tuple[int, str, str], tuple[tuple[int, str, str], ...], tuple[tuple[int, str, str], ...] +]: + """Compile the custom-rule project cold, then after each edit beside a cache-free oracle.""" + + project_dir: Path = root / "orders" + write_project_files(project_dir=project_dir, files=CUSTOM_RULES_EDIT_PROJECT) + cold: tuple[int, str, str] = rules_compile_outcome(project_dir, environment=environment) + edited: list[tuple[int, str, str]] = [] + oracles: list[tuple[int, str, str]] = [] + for step, edit_case in enumerate(edits): + edit_case.edit(project_dir) + edited.append(rules_compile_outcome(project_dir, environment=environment)) + oracle_dir: Path = root / f"oracle_{step}" + shutil.copytree(project_dir, oracle_dir, ignore=shutil.ignore_patterns("target")) + oracles.append(rules_compile_outcome(oracle_dir, "--no-cache")) + return cold, tuple(edited), tuple(oracles) + + def write_project_files(*, project_dir: Path, files: dict[str, str]) -> None: """Write authored project files beneath one project directory.""" @@ -420,7 +536,9 @@ def write_project_files(*, project_dir: Path, files: dict[str, str]) -> None: path.write_text(contents, encoding="utf-8") -def rules_compile_outcome(project_dir: Path, *arguments: str) -> tuple[int, str, str]: +def rules_compile_outcome( + project_dir: Path, *arguments: str, environment: tuple[tuple[str, str], ...] = () +) -> tuple[int, str, str]: """Compile in a fresh process and return its exit code, diagnostics, and artifacts digest.""" result: subprocess.CompletedProcess[str] = subprocess.run( @@ -436,6 +554,7 @@ def rules_compile_outcome(project_dir: Path, *arguments: str) -> tuple[int, str, capture_output=True, text=True, check=False, + env={**os.environ, **dict(environment)}, ) payload: dict[str, Any] = json.loads(result.stdout or "{}") artifacts: Any = hashlib.sha256() diff --git a/tests/e2e/src/sqlbuild/cli/commands/main/rules/test_custom_rules_incremental_sequence.py b/tests/e2e/src/sqlbuild/cli/commands/main/rules/test_custom_rules_incremental_sequence.py new file mode 100644 index 0000000000..d0c9caa44d --- /dev/null +++ b/tests/e2e/src/sqlbuild/cli/commands/main/rules/test_custom_rules_incremental_sequence.py @@ -0,0 +1,136 @@ +"""Incremental custom Rules across split hosts must equal a cache-free compile after every edit.""" + +from pathlib import Path + +import pytest + +from tests.e2e.src.sqlbuild.cli.commands.main.compile.helpers import ( + replace_project_text, + write_project_file, +) +from tests.e2e.src.sqlbuild.cli.commands.main.rules._test_types import ( + BrokenInvalidationChainCase, + RulesCacheEditCase, + RulesEditChainCase, +) +from tests.e2e.src.sqlbuild.cli.commands.main.rules.helpers import ( + STALE_SOURCE_FACTS_SITECUSTOMIZE, + custom_rules_edit_chain, + rule_cache_counts, + rules_compile_outcome, +) + +CONFIG: str = "sqlbuild_project.toml" +ORDER_TOTALS: str = "models/marts/order_totals.sql" +SOURCES: str = "sources/raw.yml" +APPROVED: str = "rules/approved_sources.yaml" +CUSTOM_EDIT_SEQUENCE: tuple[RulesCacheEditCase, ...] = ( + RulesCacheEditCase( + "model_sql_edited", + lambda root: replace_project_text(root, ORDER_TOTALS, "amount > 7", "amount > 9"), + ), + RulesCacheEditCase( + "source_in_shared_file_renamed", + lambda root: replace_project_text(root, SOURCES, "raw_products", "raw_returns"), + ), + RulesCacheEditCase( + "tracked_rule_input_edited", + lambda root: replace_project_text( + root, APPROVED, "raw_customers\n", "raw_customers\nraw_returns\n" + ), + ), + RulesCacheEditCase( + "mart_added", + lambda root: write_project_file( + root, + "models/marts/customer_totals.sql", + 'MODEL (description "Customer totals");\n\n' + 'SELECT customer_id FROM __ref("stg_customers")\n', + ), + ), + RulesCacheEditCase( + "model_removed", + lambda root: (root / "models/intermediate/orders_007.sql").unlink(), + ), + RulesCacheEditCase( + "rules_cache_disabled", + lambda root: replace_project_text( + root, + CONFIG, + "[rules.thresholds]", + "[rules.cache]\nenabled = false\n\n[rules.thresholds]", + ), + ), + RulesCacheEditCase( + "source_edited_without_cache", + lambda root: replace_project_text(root, SOURCES, "raw_returns", "raw_refunds"), + ), + RulesCacheEditCase( + "rules_cache_enabled", + lambda root: replace_project_text(root, CONFIG, "[rules.cache]\nenabled = false\n\n", ""), + ), + RulesCacheEditCase( + "source_restored_with_cache", + lambda root: replace_project_text(root, SOURCES, "raw_refunds", "raw_returns"), + ), +) + + +@pytest.mark.parametrize( + "test_case", + ( + RulesEditChainCase( + description="model_source_input_and_cache_toggle_edits", + edits=CUSTOM_EDIT_SEQUENCE, + expected_cold_codes=("XSQBRORD102", "XSQBRORD103", "XSQBRORD201"), + ), + ), + ids=lambda case: case.description, +) +def test_given_warm_custom_rules_when_applying_edit_chain_then_every_compile_equals_uncached( + tmp_path: Path, test_case: RulesEditChainCase +) -> None: + cold, edited, oracles = custom_rules_edit_chain(root=tmp_path, edits=test_case.edits) + rewarmed: tuple[int, str, str] = rules_compile_outcome(tmp_path / "orders") + + assert cold[0] == test_case.expected_exit_code + assert all(f'"{code}"' in cold[1] for code in test_case.expected_cold_codes), cold[1] + assert edited == oracles + assert {outcome[0] for outcome in edited} == {test_case.expected_exit_code} + assert '"XSQBRORD101"' in edited[0][1] + assert ("raw_returns" in edited[1][1], "raw_returns" in edited[2][1]) == (True, False) + assert rewarmed == oracles[-1] + assert rule_cache_counts(tmp_path / "orders")[1] == 0 + + +@pytest.mark.parametrize( + "test_case", + ( + BrokenInvalidationChainCase( + description="stale_shared_source_digests", + edits=CUSTOM_EDIT_SEQUENCE, + stale_digests_sitecustomize=STALE_SOURCE_FACTS_SITECUSTOMIZE, + expected_first_divergent_edit="source_in_shared_file_renamed", + ), + ), + ids=lambda case: case.description, +) +def test_given_broken_invalidation_when_applying_edit_chain_then_oracle_detects_divergence( + tmp_path: Path, test_case: BrokenInvalidationChainCase +) -> None: + write_project_file(tmp_path / "site", "sitecustomize.py", test_case.stale_digests_sitecustomize) + + _, edited, oracles = custom_rules_edit_chain( + root=tmp_path, edits=test_case.edits, environment=(("PYTHONPATH", str(tmp_path / "site")),) + ) + matches: list[bool] = [ + outcome == oracle for outcome, oracle in zip(edited, oracles, strict=True) + ] + + assert test_case.edits[matches.index(False)].description == ( + test_case.expected_first_divergent_edit + ) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/tests/unit/src/sqlbuild/rule_engine/classes/shared_fact_encoder/_test_types.py b/tests/unit/src/sqlbuild/rule_engine/classes/shared_fact_encoder/_test_types.py new file mode 100644 index 0000000000..8e009b8338 --- /dev/null +++ b/tests/unit/src/sqlbuild/rule_engine/classes/shared_fact_encoder/_test_types.py @@ -0,0 +1,16 @@ +from dataclasses import dataclass + + +@dataclass(frozen=True) +class SourceDigestTestCase: + description: str + edited_source_names: tuple[str, ...] + edited_contents: str + expected_equal: bool + + +@dataclass(frozen=True) +class ModelFactDigestTestCase: + description: str + fact: str + expected_host_digest_match: bool diff --git a/tests/unit/src/sqlbuild/rule_engine/classes/shared_fact_encoder/helpers.py b/tests/unit/src/sqlbuild/rule_engine/classes/shared_fact_encoder/helpers.py new file mode 100644 index 0000000000..1fd0d8b87b --- /dev/null +++ b/tests/unit/src/sqlbuild/rule_engine/classes/shared_fact_encoder/helpers.py @@ -0,0 +1,60 @@ +from dataclasses import replace +from pathlib import Path + +from sqlbuild.compiler.compile.models import CompiledObjectKey, CompiledProject, CompiledSource +from sqlbuild.compiler.compile.types import CompiledResourceType +from sqlbuild.compiler.discovery.models import DiscoveredSourceFile +from sqlbuild.rule_engine.classes.rule_context import RuleFactViews, build_rule_fact_views +from sqlbuild.spec.contracts.models import SourceEntry +from tests.unit.src.sqlbuild.rule_engine.main.evaluate.helpers import build_project + +SOURCES_CONTENTS: str = "sources:\n - name: raw_orders\n - name: raw_customers\n" +SOURCE_NAMES: tuple[str, ...] = ("raw_orders", "raw_customers") +_SOURCES_PATH: Path = Path("sources/raw.yml") + + +def shared_file_views(*, tmp_path: Path, names: tuple[str, ...], contents: str) -> RuleFactViews: + """Return Rule fact views over sources that share one declaration-file object.""" + + declared: DiscoveredSourceFile = _source_file(tmp_path=tmp_path, names=names, contents=contents) + return _views(tmp_path=tmp_path, files=tuple(declared for _ in names), names=names) + + +def copied_file_views(*, tmp_path: Path, names: tuple[str, ...], contents: str) -> RuleFactViews: + """Return Rule fact views over sources that each hold an equal copy of their file.""" + + files: tuple[DiscoveredSourceFile, ...] = tuple( + _source_file(tmp_path=tmp_path, names=names, contents=contents) for _ in names + ) + return _views(tmp_path=tmp_path, files=files, names=names) + + +def _source_file(*, tmp_path: Path, names: tuple[str, ...], contents: str) -> DiscoveredSourceFile: + return DiscoveredSourceFile( + file_path=tmp_path / _SOURCES_PATH, + relative_path=_SOURCES_PATH, + contents=contents, + source_entries=tuple(SourceEntry(name=name) for name in names), + ) + + +def _views( + *, tmp_path: Path, files: tuple[DiscoveredSourceFile, ...], names: tuple[str, ...] +) -> RuleFactViews: + sources: tuple[CompiledSource, ...] = tuple( + CompiledSource( + key=CompiledObjectKey(resource_type=CompiledResourceType.SOURCE, name=name), + deps=(), + name=name, + source_entry=SourceEntry(name=name), + source_file=declared, + ) + for name, declared in zip(names, files, strict=True) + ) + project: CompiledProject = replace( + build_project( + name="orders", relative_path="models/orders.sql", sql="SELECT 1", config_values={} + ), + sources=sources, + ) + return build_rule_fact_views(project=project, project_dir=tmp_path, dialect="duckdb") diff --git a/tests/unit/src/sqlbuild/rule_engine/classes/shared_fact_encoder/test_shared_fact_encoder.py b/tests/unit/src/sqlbuild/rule_engine/classes/shared_fact_encoder/test_shared_fact_encoder.py new file mode 100644 index 0000000000..44aac62117 --- /dev/null +++ b/tests/unit/src/sqlbuild/rule_engine/classes/shared_fact_encoder/test_shared_fact_encoder.py @@ -0,0 +1,108 @@ +"""Whole-project facts digest by content, however often resources share compiler objects.""" + +from pathlib import Path + +import pytest + +from sqlbuild.rule_engine._helpers.engine.fact_replay import _fact_value as fact_value +from sqlbuild.rule_engine._helpers.engine.fact_replay import ( + fact_outcome_digest, + fact_value_digest, + shared_fact_encoder, +) +from sqlbuild.rule_engine.classes.rule_context import RuleFactViews +from sqlbuild.rule_engine.constants import ( + FACT_COLUMNS_DECLARED, + FACT_PROJECT_SOURCES, + FACT_SQL_FOR_MODEL, +) +from sqlbuild.rule_engine.types import FactKey +from tests.unit.src.sqlbuild.rule_engine.classes.shared_fact_encoder._test_types import ( + ModelFactDigestTestCase, + SourceDigestTestCase, +) +from tests.unit.src.sqlbuild.rule_engine.classes.shared_fact_encoder.helpers import ( + SOURCE_NAMES, + SOURCES_CONTENTS, + copied_file_views, + shared_file_views, +) + + +@pytest.mark.parametrize( + "test_case", + ( + SourceDigestTestCase( + description="unchanged declarations", + edited_source_names=SOURCE_NAMES, + edited_contents=SOURCES_CONTENTS, + expected_equal=True, + ), + SourceDigestTestCase( + description="shared file contents edited", + edited_source_names=SOURCE_NAMES, + edited_contents=SOURCES_CONTENTS + "# reviewed\n", + expected_equal=False, + ), + SourceDigestTestCase( + description="source renamed", + edited_source_names=("raw_orders", "raw_clients"), + edited_contents=SOURCES_CONTENTS, + expected_equal=False, + ), + SourceDigestTestCase( + description="source added", + edited_source_names=(*SOURCE_NAMES, "raw_products"), + edited_contents=SOURCES_CONTENTS, + expected_equal=False, + ), + ), + ids=lambda case: case.description, +) +def test_given_sources_sharing_a_file_when_digesting_then_digest_follows_content_only( + tmp_path: Path, test_case: SourceDigestTestCase +) -> None: + shared: RuleFactViews = shared_file_views( + tmp_path=tmp_path, names=SOURCE_NAMES, contents=SOURCES_CONTENTS + ) + copied: RuleFactViews = copied_file_views( + tmp_path=tmp_path, + names=test_case.edited_source_names, + contents=test_case.edited_contents, + ) + + before: str = fact_outcome_digest( + views=shared, key=(FACT_PROJECT_SOURCES,), shared=shared_fact_encoder() + ) + after: str = fact_outcome_digest( + views=copied, key=(FACT_PROJECT_SOURCES,), shared=shared_fact_encoder() + ) + + assert (before == after) is test_case.expected_equal + + +@pytest.mark.parametrize( + "test_case", + ( + ModelFactDigestTestCase("expanded and authored SQL", FACT_SQL_FOR_MODEL, True), + ModelFactDigestTestCase("declared columns", FACT_COLUMNS_DECLARED, True), + ), + ids=lambda case: case.description, +) +def test_given_model_fact_when_digesting_with_shared_encoder_then_matches_host_observed_digest( + tmp_path: Path, test_case: ModelFactDigestTestCase +) -> None: + views: RuleFactViews = shared_file_views( + tmp_path=tmp_path, names=SOURCE_NAMES, contents=SOURCES_CONTENTS + ) + key: FactKey = (test_case.fact, views.project.models[0].path.as_posix()) + + digest: str = fact_outcome_digest(views=views, key=key, shared=shared_fact_encoder()) + + assert (digest == fact_value_digest(fact_value(views=views, key=key))) is ( + test_case.expected_host_digest_match + ) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-vv"]))