diff --git a/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs b/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs index fdc5c9b548..59b33f5732 100644 --- a/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs +++ b/crates/sqlbuild-rules-native/src/bindings/_helpers/functions.rs @@ -117,6 +117,31 @@ fn prepare_lint_sql( .map_err(value_error) } +/// Prepare many bodies in one detached call; `false` marks a body to prepare again individually. +#[pyfunction] +fn prepare_lint_sql_batch( + py: Python<'_>, + requests: Vec, +) -> PyResult)>> { + py.compiler_detach(|| { + Ok(requests + .iter() + .map(|request| { + crate::bindings::_helpers::panics::catch_compiler_panic(|| { + crate::sql_lint::main::preparation::prepare( + &request.expanded, + &request.before_expansion, + &request.prior_sites, + &request.dialect, + ) + }) + .map_or((false, None), |prepared| (true, prepared)) + }) + .collect()) + }) + .map_err(value_error) +} + #[pyfunction] fn lint_backtick_identifiers(dialect: &str) -> PyResult { compiler_guard(|| { @@ -431,6 +456,7 @@ pub(crate) fn register(module: &Bound<'_, PyModule>) -> PyResult<()> { module.add_function(wrap_pyfunction!(run_custom_host_json, module)?)?; module.add_function(wrap_pyfunction!(lint_sql_json, module)?)?; module.add_function(wrap_pyfunction!(prepare_lint_sql, module)?)?; + module.add_function(wrap_pyfunction!(prepare_lint_sql_batch, module)?)?; module.add( "NativeCompilerError", module diff --git a/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/rules/custom-rules/testing-and-determinism.md b/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/rules/custom-rules/testing-and-determinism.md index c3641a48d4..ef11cb0bf5 100644 --- a/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/rules/custom-rules/testing-and-determinism.md +++ b/src/sqlbuild/.agents/skills/sqlbuild/references/docs/concepts/rules/custom-rules/testing-and-determinism.md @@ -6,6 +6,14 @@ Online: https://sqlbuild.com/docs/concepts/rules/custom-rules/testing-and-determinism/ +## Contents + +- Test Rules +- Parallel evaluation +- Helpers +- Deterministic inputs +- Cache granularity + ## Test Rules Use the public harness to exercise discovery, compilation, and Rules evaluation: @@ -39,6 +47,19 @@ the next. Compute findings from the subject and `ctx` alone. Repository pytest files remain outside the SQLBuild project's SQL `tests/` directory. +## Parallel evaluation + +On large projects, SQLBuild splits custom Rule evaluation across several host processes, at most +one per available CPU and never more than eight. Each process evaluates a share of the subjects, +and SQLBuild merges the findings back into the order a single process reports them. + +A Rule that changes module-level state while it runs, for example by adding each model to a set of +already-seen names, could report different findings when its subjects are split. SQLBuild detects +the state change and evaluates that Rule again in one process over all of its subjects, so its +findings and cache entries match a single-process run. The result stays correct, but the Rule no +longer runs in parallel and costs an extra host process. Keep each subject's findings independent +of the others so Rules can run in parallel. + ## Helpers Custom Rules can import repository-owned helpers under `rules/` and these deterministic standard diff --git a/src/sqlbuild/_native.pyi b/src/sqlbuild/_native.pyi index 5c54c88e00..d90c492dac 100644 --- a/src/sqlbuild/_native.pyi +++ b/src/sqlbuild/_native.pyi @@ -23,6 +23,9 @@ def lint_sql_batch_json(request_json: str) -> str: ... def prepare_lint_sql( request: LintPreparationRequest, ) -> tuple[str, list[tuple[str, int, int, int, int, str]], list[str]] | None: ... +def prepare_lint_sql_batch( + requests: list[LintPreparationRequest], +) -> list[tuple[bool, tuple[str, list[tuple[str, int, int, int, int, str]], list[str]] | None]]: ... def lint_backtick_identifiers(dialect: str) -> bool: ... def load_config_json(project_dir: str) -> str: ... def selected_codes_json(request_json: str) -> str: ... diff --git a/src/sqlbuild/lint/_helpers/expansion.py b/src/sqlbuild/lint/_helpers/expansion.py index 5b641ee845..29ae933e18 100644 --- a/src/sqlbuild/lint/_helpers/expansion.py +++ b/src/sqlbuild/lint/_helpers/expansion.py @@ -3,6 +3,7 @@ from __future__ import annotations import re +from dataclasses import replace from pathlib import Path import sqlbuild._native as _native @@ -23,15 +24,32 @@ from sqlbuild.compiler.discovery.main.discover import discover_project_inputs from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs from sqlbuild.compiler.scopes.types import DeclarationKind +from sqlbuild.lint._helpers.headers import lint_body_ranges, lint_file_role +from sqlbuild.lint._helpers.native import external_identifiers_for_headers +from sqlbuild.lint._helpers.native_sql import check_lint_not_stopped from sqlbuild.lint._helpers.sqlbuild_tokens import ( neutralize_context_interpolation, neutralize_generic_audit_parameters, neutralize_interpolation, sentinel_spans, ) -from sqlbuild.lint.constants import RELATION_IDENTITY_TEMPLATE, TEMPLATE_INTERPOLATION_START +from sqlbuild.lint.constants import ( + HEADER_KIND_SCENARIO, + HEADER_KIND_TEST, + RELATION_IDENTITY_TEMPLATE, + TEMPLATE_INTERPOLATION_START, +) from sqlbuild.lint.exceptions import ProjectCompileError -from sqlbuild.lint.models import InterpolationSite, LintBody, LintFileRole +from sqlbuild.lint.models import ( + ExpandedLintFile, + HeaderSpan, + InterpolationSite, + LintBody, + LintFileRole, + PendingLintBody, + PreparedLintFiles, +) +from sqlbuild.lint.types import NativePreparedSql from sqlbuild.spec.contracts.main.resolve_effective_adapter_name import ( resolve_effective_adapter_name, ) @@ -146,6 +164,38 @@ def prepare_lint_body( ) -> LintBody: """Expand one authored body and neutralize whatever interpolation remains.""" + pending: PendingLintBody = expand_lint_body( + role=role, + file_path=file_path, + contents=contents, + body_range=body_range, + context=context, + dialect=dialect, + external_identifiers=external_identifiers, + allows_ceremonial_select=allows_ceremonial_select, + allows_dynamic_output_star=allows_dynamic_output_star, + compiled_expansion=compiled_expansion, + ) + return finish_lint_body( + pending=pending, prepared=_native.prepare_lint_sql(pending.request), dialect=dialect + ) + + +def expand_lint_body( + *, + role: LintFileRole, + file_path: Path, + contents: str, + body_range: tuple[int, int], + context: SqlExpansionContext, + dialect: str, + external_identifiers: tuple[str, ...] = (), + allows_ceremonial_select: bool = False, + allows_dynamic_output_star: bool = False, + compiled_expansion: CompiledSqlExpansion | None = None, +) -> PendingLintBody: + """Expand one authored body like compile does, leaving its native preparation pending.""" + body_start: int body_end: int body_start, body_end = body_range @@ -184,50 +234,84 @@ def prepare_lint_body( raise ProjectCompileError( f"{file_path} could not be expanded, so its SQL cannot be linted: {error}" ) from error + return PendingLintBody( + file_path=file_path, + body_start=body_start, + body_end=body_end, + expansion_input=expansion_input, + expanded=expanded, + expansion_passes=expansion_passes, + pre_expansion_sites=pre_expansion_sites, + external_identifiers=external_identifiers, + allows_ceremonial_select=allows_ceremonial_select, + allows_dynamic_output_star=allows_dynamic_output_star, + allows_empty_fixture_star=False, + request={ + "expanded": expanded, + "before_expansion": expansion_input, + "prior_sites": [site.neutralized_start for site in pre_expansion_sites], + "dialect": dialect, + }, + ) + + +def prepare_lint_sql_batch( + pending: tuple[PendingLintBody, ...], +) -> tuple[NativePreparedSql | None, ...]: + """Prepare pending bodies in one GIL-free native call; failures re-raise per body, in order.""" + + if not pending: + return () + results: list[tuple[bool, NativePreparedSql | None]] = _native.prepare_lint_sql_batch( + [body.request for body in pending] + ) + return tuple( + prepared if succeeded else _native.prepare_lint_sql(body.request) + for body, (succeeded, prepared) in zip(pending, results, strict=True) + ) + + +def finish_lint_body( + *, pending: PendingLintBody, prepared: NativePreparedSql | None, dialect: str +) -> LintBody: + """Build the lint body from one body's native preparation, or the Python fallback.""" + + expanded: str = pending.expanded neutralized: str sites: tuple[InterpolationSite, ...] - native_prepared: tuple[str, list[tuple[str, int, int, int, int, str]], list[str]] | None = ( - _native.prepare_lint_sql( - { - "expanded": expanded, - "before_expansion": expansion_input, - "prior_sites": [site.neutralized_start for site in pre_expansion_sites], - "dialect": dialect, - } - ) - ) - if native_prepared is None: + if prepared is None: neutralized, sites = neutralize_interpolation(body=expanded, dialect=dialect) externally_referenced_ctes: tuple[str, ...] = _externally_referenced_ctes( expanded=expanded, interpolation_sites=sites, - pre_expansion_body=expansion_input, - pre_expansion_sites=pre_expansion_sites, + pre_expansion_body=pending.expansion_input, + pre_expansion_sites=pending.pre_expansion_sites, ) else: - neutralized = native_prepared[0] - sites = tuple(InterpolationSite(*site) for site in native_prepared[1]) - externally_referenced_ctes = tuple(native_prepared[2]) + neutralized = prepared[0] + sites = tuple(InterpolationSite(*site) for site in prepared[1]) + externally_referenced_ctes = tuple(prepared[2]) dependency_identifiers: tuple[str, ...] = tuple( site.sentinel for site in sites if _DEPENDENCY_INTRINSIC_PATTERN.match(site.original_text) is not None ) return LintBody( - file_path=file_path, - body_start=body_start, - body_end=body_end, + file_path=pending.file_path, + body_start=pending.body_start, + body_end=pending.body_end, lint_text=neutralized, passes=( - sentinel_spans(sites=pre_expansion_sites), - *expansion_passes, + sentinel_spans(sites=pending.pre_expansion_sites), + *pending.expansion_passes, sentinel_spans(sites=sites), ), - external_identifiers=external_identifiers, + external_identifiers=pending.external_identifiers, dependency_identifiers=dependency_identifiers, externally_referenced_ctes=externally_referenced_ctes, - allows_ceremonial_select=allows_ceremonial_select, - allows_dynamic_output_star=allows_dynamic_output_star, + allows_ceremonial_select=pending.allows_ceremonial_select, + allows_dynamic_output_star=pending.allows_dynamic_output_star, + allows_empty_fixture_star=pending.allows_empty_fixture_star, dependency_relations=tuple( ( site.sentinel, @@ -249,3 +333,96 @@ def _is_sql_hook_path(*, role: LintFileRole) -> bool: return role.in_hook_directory or ( role.in_project and role.declaration_kind is DeclarationKind.SQL_HOOK ) + + +def expand_file_bodies( + *, + file_path: Path, + contents: str, + headers: tuple[HeaderSpan, ...], + context: SqlExpansionContext, + dialect: str, + project_dir: Path, + relative_path: Path, + dynamic_output_paths: frozenset[Path], + compiled_expansions: dict[Path, CompiledSqlExpansion] | None, +) -> ExpandedLintFile: + """Expand one file's bodies up to the first body that cannot be expanded.""" + + check_lint_not_stopped() + bodies: list[PendingLintBody] = [] + role: LintFileRole = lint_file_role( + file_path=file_path, project_dir=project_dir, relative_path=relative_path + ) + allows_dynamic_output_star: bool = ( + bool(dynamic_output_paths) and file_path.resolve() in dynamic_output_paths + ) + compiled_expansion: CompiledSqlExpansion | None = (compiled_expansions or {}).get(file_path) + external_identifiers: tuple[str, ...] = external_identifiers_for_headers( + contents=contents, headers=headers + ) + allows_ceremonial_select: bool = any( + header.kind in {HEADER_KIND_TEST, HEADER_KIND_SCENARIO} for header in headers + ) + allows_empty_fixture_star: bool = any(header.kind == HEADER_KIND_TEST for header in headers) + body_start: int + body_end: int + try: + for body_start, body_end in lint_body_ranges( + contents=contents, + headers=headers, + file_path=file_path, + project_dir=project_dir, + role=role, + ): + bodies.append( + replace( + expand_lint_body( + role=role, + file_path=file_path, + contents=contents, + body_range=(body_start, body_end), + context=context, + dialect=dialect, + external_identifiers=external_identifiers, + allows_ceremonial_select=allows_ceremonial_select, + allows_dynamic_output_star=allows_dynamic_output_star, + compiled_expansion=compiled_expansion, + ), + allows_empty_fixture_star=allows_empty_fixture_star, + ) + ) + except ProjectCompileError as error: + return ExpandedLintFile(file_path=file_path, bodies=tuple(bodies), failure=error) + return ExpandedLintFile(file_path=file_path, bodies=tuple(bodies), failure=None) + + +def prepare_expanded_files( + *, expanded_files: tuple[ExpandedLintFile, ...], dialect: str, skip_unexpandable: bool +) -> PreparedLintFiles: + """Prepare every expanded body natively at once, then report failures in file order.""" + + check_lint_not_stopped() + pending: list[PendingLintBody] = [] + for expanded in expanded_files: + pending.extend(expanded.bodies) + prepared: tuple[NativePreparedSql | None, ...] = prepare_lint_sql_batch(tuple(pending)) + bodies: list[LintBody] = [] + unexpandable: dict[Path, str] = {} + position: int = 0 + for expanded in expanded_files: + file_bodies: list[LintBody] = [] + for body in expanded.bodies: + file_bodies.append( + finish_lint_body(pending=body, prepared=prepared[position], dialect=dialect) + ) + position += 1 + if expanded.failure is None: + bodies.extend(file_bodies) + elif skip_unexpandable: + unexpandable[expanded.file_path] = str(expanded.failure).removeprefix( + f"{expanded.file_path} " + ) + else: + raise expanded.failure + return PreparedLintFiles(bodies=tuple(bodies), unexpandable=unexpandable) diff --git a/src/sqlbuild/lint/main/run_lint.py b/src/sqlbuild/lint/main/run_lint.py index f746cd5975..9756337ee1 100644 --- a/src/sqlbuild/lint/main/run_lint.py +++ b/src/sqlbuild/lint/main/run_lint.py @@ -2,7 +2,6 @@ from __future__ import annotations -from dataclasses import replace from pathlib import Path from sqlbuild.compiler.compile.constants import MODEL_DIRECTORY_NAME @@ -14,21 +13,24 @@ INHERITED_DECLARATION_DIRECTORIES, LOCAL_DECLARATION_DIRECTORIES, ) -from sqlbuild.lint._helpers.expansion import build_lint_expansion_context, prepare_lint_body -from sqlbuild.lint._helpers.headers import lint_body_ranges, lint_file_role, scan_headers -from sqlbuild.lint._helpers.native import external_identifiers_for_headers, lint_native_headers -from sqlbuild.lint._helpers.native_sql import check_lint_not_stopped, run_native_sql_lint +from sqlbuild.lint._helpers.expansion import ( + build_lint_expansion_context, + expand_file_bodies, + prepare_expanded_files, +) +from sqlbuild.lint._helpers.headers import scan_headers +from sqlbuild.lint._helpers.native import lint_native_headers +from sqlbuild.lint._helpers.native_sql import run_native_sql_lint from sqlbuild.lint._helpers.project_files import collect_project_files, sort_violations from sqlbuild.lint._helpers.suppressions import apply_suppressions -from sqlbuild.lint.constants import HEADER_KIND_SCENARIO, HEADER_KIND_TEST -from sqlbuild.lint.exceptions import ProjectCompileError from sqlbuild.lint.models import ( + ExpandedLintFile, HeaderSpan, LintBody, LintConfig, - LintFileRole, LintRunResult, LintViolation, + PreparedLintFiles, ) @@ -51,62 +53,51 @@ def run_lint( project_dir=project_dir, selected_paths=selected_paths, source_files=source_files ) violations: list[LintViolation] = [] - bodies: list[LintBody] = [] - unexpandable: dict[Path, str] = {} context: SqlExpansionContext | None = expansion_context or _expansion_context( project_dir=project_dir, native_enabled=config.native_enabled, value_renderer=value_renderer, discovered_inputs=discovered_inputs, ) + expanded_files: list[ExpandedLintFile] = [] + failure: Exception | None = None file_path: Path contents: str for file_path, contents in sorted(files.items()): - relative_path: Path = file_path.relative_to(project_dir) - relative_parts: tuple[str, ...] = relative_path.parts - declaration_directories: frozenset[str] = ( - INHERITED_DECLARATION_DIRECTORIES - | LOCAL_DECLARATION_DIRECTORIES - | {DECLARATION_GROUP_DIRECTORY} - ) - headers: tuple[HeaderSpan, ...] = scan_headers( - contents=contents, - first_only=( - relative_parts[:1] == (MODEL_DIRECTORY_NAME,) - and not declaration_directories.intersection(relative_parts[1:-1]) - ), - ) - if config.header_rules_enabled: - violations.extend( - lint_native_headers( - contents=contents, - file_path=file_path, - headers=headers, - config=config, - ) + try: + header_violations: tuple[LintViolation, ...] + expanded: ExpandedLintFile | None + header_violations, expanded = _scan_file( + file_path=file_path, + contents=contents, + project_dir=project_dir, + config=config, + context=context, + dynamic_output_paths=dynamic_output_paths, + compiled_expansions=compiled_expansions, ) - if context is None: + except Exception as error: + failure = error + break + violations.extend(header_violations) + if expanded is None: continue - prepared: tuple[LintBody, ...] | str = _prepared_bodies( - file_path=file_path, - contents=contents, - headers=headers, - context=context, - dialect=config.dialect, - project_dir=project_dir, - relative_path=relative_path, - dynamic_output_paths=dynamic_output_paths, - compiled_expansions=compiled_expansions, - skip_unexpandable=skip_unexpandable, - ) - if isinstance(prepared, str): - unexpandable[file_path] = prepared - else: - bodies.extend(prepared) + expanded_files.append(expanded) + if expanded.failure is not None and not skip_unexpandable: + break + prepared: PreparedLintFiles = prepare_expanded_files( + expanded_files=tuple(expanded_files), + dialect=config.dialect, + skip_unexpandable=skip_unexpandable, + ) + if failure is not None: + raise failure + bodies: tuple[LintBody, ...] = prepared.bodies + unexpandable: dict[Path, str] = prepared.unexpandable if bodies and config.native_enabled: native_violations: dict[Path, tuple[LintViolation, ...]] = run_native_sql_lint( - bodies=tuple(bodies), + bodies=bodies, contents_by_path=files, config=config, ) @@ -139,66 +130,54 @@ def _expansion_context( ) -def _prepared_bodies( +def _scan_file( *, file_path: Path, contents: str, - headers: tuple[HeaderSpan, ...], - context: SqlExpansionContext, - dialect: str, project_dir: Path, - relative_path: Path, + config: LintConfig, + context: SqlExpansionContext | None, dynamic_output_paths: frozenset[Path], compiled_expansions: dict[Path, CompiledSqlExpansion] | None, - skip_unexpandable: bool, -) -> tuple[LintBody, ...] | str: - """Prepared bodies of one file, or the reason its SQL cannot be expanded when skipping.""" +) -> tuple[tuple[LintViolation, ...], ExpandedLintFile | None]: + """Header violations of one file and, when expanding, its bodies awaiting preparation.""" - check_lint_not_stopped() - bodies: list[LintBody] = [] - role: LintFileRole = lint_file_role( - file_path=file_path, project_dir=project_dir, relative_path=relative_path + relative_path: Path = file_path.relative_to(project_dir) + relative_parts: tuple[str, ...] = relative_path.parts + declaration_directories: frozenset[str] = ( + INHERITED_DECLARATION_DIRECTORIES + | LOCAL_DECLARATION_DIRECTORIES + | {DECLARATION_GROUP_DIRECTORY} ) - allows_dynamic_output_star: bool = ( - bool(dynamic_output_paths) and file_path.resolve() in dynamic_output_paths + headers: tuple[HeaderSpan, ...] = scan_headers( + contents=contents, + first_only=( + relative_parts[:1] == (MODEL_DIRECTORY_NAME,) + and not declaration_directories.intersection(relative_parts[1:-1]) + ), ) - compiled_expansion: CompiledSqlExpansion | None = (compiled_expansions or {}).get(file_path) - external_identifiers: tuple[str, ...] = external_identifiers_for_headers( - contents=contents, headers=headers + header_violations: tuple[LintViolation, ...] = ( + tuple( + lint_native_headers( + contents=contents, + file_path=file_path, + headers=headers, + config=config, + ) + ) + if config.header_rules_enabled + else () ) - allows_ceremonial_select: bool = any( - header.kind in {HEADER_KIND_TEST, HEADER_KIND_SCENARIO} for header in headers + if context is None: + return header_violations, None + return header_violations, expand_file_bodies( + file_path=file_path, + contents=contents, + headers=headers, + context=context, + dialect=config.dialect, + project_dir=project_dir, + relative_path=relative_path, + dynamic_output_paths=dynamic_output_paths, + compiled_expansions=compiled_expansions, ) - allows_empty_fixture_star: bool = any(header.kind == HEADER_KIND_TEST for header in headers) - body_start: int - body_end: int - try: - for body_start, body_end in lint_body_ranges( - contents=contents, - headers=headers, - file_path=file_path, - project_dir=project_dir, - role=role, - ): - bodies.append( - replace( - prepare_lint_body( - role=role, - file_path=file_path, - contents=contents, - body_range=(body_start, body_end), - context=context, - dialect=dialect, - external_identifiers=external_identifiers, - allows_ceremonial_select=allows_ceremonial_select, - allows_dynamic_output_star=allows_dynamic_output_star, - compiled_expansion=compiled_expansion, - ), - allows_empty_fixture_star=allows_empty_fixture_star, - ) - ) - except ProjectCompileError as error: - if not skip_unexpandable: - raise - return str(error).removeprefix(f"{file_path} ") - return tuple(bodies) diff --git a/src/sqlbuild/lint/models.py b/src/sqlbuild/lint/models.py index 296ac9ca06..eb88a403d0 100644 --- a/src/sqlbuild/lint/models.py +++ b/src/sqlbuild/lint/models.py @@ -15,7 +15,13 @@ VIOLATION_SEVERITY_FAULT, VIOLATION_SEVERITY_WARNING, ) -from sqlbuild.lint.types import DiagnosticIdentity, LintSeverity, RuleFixStatus +from sqlbuild.lint.exceptions import ProjectCompileError +from sqlbuild.lint.types import ( + DiagnosticIdentity, + LintSeverity, + NativeLintPreparationRequest, + RuleFixStatus, +) @dataclass(frozen=True) @@ -66,6 +72,41 @@ class LintBody: dependency_relations: tuple[tuple[str, str], ...] = () +@dataclass(frozen=True) +class PendingLintBody: + """One expanded authored body awaiting its native lexical preparation.""" + + file_path: Path + body_start: int + body_end: int + expansion_input: str + expanded: str + expansion_passes: tuple[tuple[ExpansionSpan, ...], ...] + pre_expansion_sites: tuple[InterpolationSite, ...] + external_identifiers: tuple[str, ...] + allows_ceremonial_select: bool + allows_dynamic_output_star: bool + allows_empty_fixture_star: bool + request: NativeLintPreparationRequest + + +@dataclass(frozen=True) +class ExpandedLintFile: + """The expanded bodies of one file, up to the first body that could not be expanded.""" + + file_path: Path + bodies: tuple[PendingLintBody, ...] + failure: ProjectCompileError | None + + +@dataclass(frozen=True) +class PreparedLintFiles: + """Lint bodies of every expandable file, and why the skipped files could not expand.""" + + bodies: tuple[LintBody, ...] + unexpandable: dict[Path, str] + + @dataclass(frozen=True) class FixtureNullCandidateScan: """Cheaply discovered SQL-test fixtures that may contain removable typed nulls.""" diff --git a/src/sqlbuild/lint/types.py b/src/sqlbuild/lint/types.py index 6909aadd68..82135aafcf 100644 --- a/src/sqlbuild/lint/types.py +++ b/src/sqlbuild/lint/types.py @@ -3,7 +3,10 @@ from __future__ import annotations from enum import StrEnum -from typing import Literal +from typing import TYPE_CHECKING, Literal + +if TYPE_CHECKING: + from sqlbuild._native import LintPreparationRequest class LintSeverity(StrEnum): @@ -18,3 +21,7 @@ class LintSeverity(StrEnum): type RelationKeys = dict[str, tuple[tuple[str, ...], ...]] type DiagnosticIdentity = tuple[str, str, str, str] + +type NativeLintPreparationRequest = LintPreparationRequest + +type NativePreparedSql = tuple[str, list[tuple[str, int, int, int, int, str]], list[str]] diff --git a/src/sqlbuild/rule_engine/_helpers/engine/catalogue.py b/src/sqlbuild/rule_engine/_helpers/engine/catalogue.py index 626acd3757..1ece17e899 100644 --- a/src/sqlbuild/rule_engine/_helpers/engine/catalogue.py +++ b/src/sqlbuild/rule_engine/_helpers/engine/catalogue.py @@ -174,13 +174,20 @@ def _native_builtin(*, model: object, ctx: object) -> list[Finding]: def select_rules( - *, catalogue: tuple[Rule, ...], config: RulesConfig, project_dir: Path + *, + catalogue: tuple[Rule, ...], + config: RulesConfig, + project_dir: Path, + custom_payloads: list[dict[str, object]] | None = None, ) -> tuple[Rule, ...]: """Resolve the active catalogue through the native Fensu rules owner.""" by_code: dict[str, Rule] = {rule.code: rule for rule in catalogue} codes: tuple[str, ...] = native_selected_codes( - config=config, catalogue=catalogue, project_dir=project_dir + config=config, + catalogue=catalogue, + project_dir=project_dir, + custom_payloads=custom_payloads, ) try: return tuple(by_code[code] for code in codes) diff --git a/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py b/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py index 4b500db707..9833dc1e18 100644 --- a/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py +++ b/src/sqlbuild/rule_engine/_helpers/engine/custom_rules.py @@ -4,19 +4,15 @@ import hashlib import os -import pickle -import sys -import tempfile import threading import time from collections.abc import Callable from dataclasses import asdict, dataclass, replace from pathlib import Path -from typing import Any, cast +from typing import Any import orjson -import sqlbuild._native as _native from sqlbuild.compiler.compile.models import CompiledProject from sqlbuild.compiler.fact_cache.main.code_identity import compiled_code_identity from sqlbuild.rule_engine._helpers.engine.custom_rule_evidence import ( @@ -24,21 +20,29 @@ custom_rule_import_closure, ) from sqlbuild.rule_engine._helpers.engine.fact_replay import full_fact_keys, is_tree_fact +from sqlbuild.rule_engine._helpers.host.custom_host_pool import ( + decode_readset, + finding_payload, + run_custom_hosts, +) from sqlbuild.rule_engine.classes.fact_digests import FactDigests from sqlbuild.rule_engine.classes.rule_context import RuleFactViews, build_rule_fact_views from sqlbuild.rule_engine.constants import ( - CUSTOM_HOST_LAUNCH_MODULE, CUSTOM_HOST_RUNTIME_VERSION, + CUSTOM_RULE_PROJECT_SUBJECT, CUSTOM_RULES_CACHE_FILE, CUSTOM_RULES_CACHE_VERSION, ) from sqlbuild.rule_engine.exceptions import FactDigestError, RulesError -from sqlbuild.rule_engine.models import CustomRulesOutcome, Rule, RulesConfig +from sqlbuild.rule_engine.models import ( + CustomHostEvaluation, + CustomHostRun, + CustomRulesOutcome, + Rule, + RulesConfig, +) from sqlbuild.rule_engine.types import FactKey, RuleSubject -_HOST_TIMEOUT_MILLIS: int = 120_000 -_PROJECT_SUBJECT: str = "" - @dataclass(frozen=True) class _Subject: @@ -46,20 +50,6 @@ class _Subject: identity: str -@dataclass(frozen=True) -class _Evaluated: - findings: tuple[dict[str, object], ...] - reads: tuple[FactKey, ...] | None - - -@dataclass(frozen=True) -class _HostRun: - evaluated: dict[tuple[str, str], _Evaluated] - untracked_reads: dict[str, tuple[FactKey, ...]] - uncacheable: frozenset[str] - observed: dict[FactKey, str | None] - - @dataclass(frozen=True) class _CachedEntry: subject_identity: str @@ -83,6 +73,7 @@ def evaluate_custom_rules_cached( rules: tuple[Rule, ...], dialect: str, verify_determinism: bool = False, + implementation_fingerprints: dict[str, str] | None = None, ) -> CustomRulesOutcome: """Reuse results whose recorded fact reads are unchanged and evaluate only the rest.""" @@ -91,11 +82,13 @@ def evaluate_custom_rules_cached( if not custom: return CustomRulesOutcome(findings=(), cache_hits=0, cache_misses=0, custom_ms=0) if not config.cache.enabled: - uncached: _HostRun = _run_host( - project=project, + uncached: CustomHostRun = run_custom_hosts( + project=_host_project(project), config=config, project_dir=project_dir, dialect=dialect, + rules=custom, + model_paths=_model_paths(project), plan={rule.code: None for rule in custom}, track_reads=False, verify_determinism=verify_determinism, @@ -116,7 +109,11 @@ def evaluate_custom_rules_cached( subjects: dict[str, tuple[_Subject, ...]] = _subjects(project=project, rules=custom) identities: dict[str, str] = { rule.code: _rule_identity( - rule=rule, config=config, project_dir=project_dir, dialect=dialect + rule=rule, + config=config, + project_dir=project_dir, + dialect=dialect, + fingerprint=(implementation_fingerprints or {}).get(rule.code), ) for rule in custom } @@ -152,11 +149,13 @@ def evaluate_custom_rules_cached( len(subjects[code]) if planned is None else len(planned) for code, planned in plan.items() ) if plan: - run: _HostRun = _run_host( - project=project, + run: CustomHostRun = run_custom_hosts( + project=_host_project(project), config=config, project_dir=project_dir, dialect=dialect, + rules=custom, + model_paths=_model_paths(project), plan=plan, track_reads=True, verify_determinism=verify_determinism, @@ -200,7 +199,7 @@ def _subjects( ) -> dict[str, tuple[_Subject, ...]]: project_subject: tuple[_Subject, ...] = ( _Subject( - key=_PROJECT_SUBJECT, + key=CUSTOM_RULE_PROJECT_SUBJECT, identity=_digest([len(project.models), project.effective_target_name]), ), ) @@ -217,13 +216,18 @@ def _subjects( } +def _model_paths(project: CompiledProject) -> tuple[str, ...]: + return tuple(sorted(model.relative_path.as_posix() for model in project.models)) + + def _materialization(model: Any) -> str | None: value: object = model.config.values.get("materialized") return value if isinstance(value, str) else None -def _rule_identity(*, rule: Rule, config: RulesConfig, project_dir: Path, dialect: str) -> str: - closure: tuple[Path, ...] = custom_rule_import_closure(rule=rule, project_dir=project_dir) +def _rule_identity( + *, rule: Rule, config: RulesConfig, project_dir: Path, dialect: str, fingerprint: str | None +) -> str: return _digest( [ CUSTOM_RULES_CACHE_VERSION, @@ -235,8 +239,12 @@ def _rule_identity(*, rule: Rule, config: RulesConfig, project_dir: Path, dialec rule.message, rule.remediation, None if rule.subject is None else rule.subject.value, - custom_rule_implementation_fingerprint( - rule=rule, project_dir=project_dir, import_closure=closure + fingerprint + if fingerprint is not None + else custom_rule_implementation_fingerprint( + rule=rule, + project_dir=project_dir, + import_closure=custom_rule_import_closure(rule=rule, project_dir=project_dir), ), asdict(config), dialect, @@ -260,51 +268,6 @@ def _readset_digest( return memo[index] -def _run_host( - *, - project: CompiledProject, - config: RulesConfig, - project_dir: Path, - dialect: str, - plan: dict[str, list[str] | None], - track_reads: bool, - verify_determinism: bool, -) -> _HostRun: - input_dir: Path = project_dir / "target" / "rules-cache" / "host-inputs" - input_dir.mkdir(parents=True, exist_ok=True) - input_path: Path | None = None - try: - with tempfile.NamedTemporaryFile( - dir=input_dir, prefix="project-", suffix=".pickle", delete=False - ) as handle: - input_path = Path(handle.name) - pickle.dump((_host_project(project), config), handle) - response_json: str = _native.run_custom_host_json( - orjson.dumps( - { - "program": sys.executable, - "arguments": ["-m", CUSTOM_HOST_LAUNCH_MODULE], - "timeout_millis": _HOST_TIMEOUT_MILLIS, - "runtime_version": CUSTOM_HOST_RUNTIME_VERSION, - "payload": { - "project_pickle_path": str(input_path.resolve()), - "project_dir": str(project_dir.resolve()), - "dialect": dialect, - "verify_determinism": verify_determinism, - "track_reads": track_reads, - "plan": plan, - }, - } - ).decode() - ) - except (ValueError, TypeError) as error: - raise RulesError(str(error)) from error - finally: - if input_path is not None: - input_path.unlink(missing_ok=True) - return _decode_host_response(orjson.loads(response_json)) - - def _host_project(project: CompiledProject) -> CompiledProject: return replace( project, @@ -318,59 +281,11 @@ def _host_project(project: CompiledProject) -> CompiledProject: ) -def _decode_host_response(response: object) -> _HostRun: - if not isinstance(response, dict): - raise RulesError("custom rule host returned an invalid payload") - payload: dict[str, Any] = cast(dict[str, Any], response) - raw_readsets: Any = payload.get("readsets") - raw_evaluations: Any = payload.get("evaluations") - raw_untracked: Any = payload.get("untracked") - raw_uncacheable: Any = payload.get("uncacheable") - raw_observed: Any = payload.get("observed") - if ( - not isinstance(raw_readsets, list) - or not isinstance(raw_evaluations, list) - or not isinstance(raw_untracked, dict) - or not isinstance(raw_uncacheable, list) - or not isinstance(raw_observed, list) - ): - raise RulesError("custom rule host returned an invalid payload") - readsets: list[tuple[FactKey, ...]] = list(map(_decode_readset, raw_readsets)) - untracked: dict[str, tuple[FactKey, ...]] = { - str(code): _decode_readset(reads) for code, reads in raw_untracked.items() - } - uncacheable: frozenset[str] = frozenset(str(code) for code in raw_uncacheable) - evaluated: dict[tuple[str, str], _Evaluated] = {} - for raw_item in raw_evaluations: - if not isinstance(raw_item, dict): - raise RulesError("custom rule host returned an invalid evaluation") - item: dict[str, Any] = cast(dict[str, Any], raw_item) - code: str = str(item["code"]) - subject: object = item.get("subject") - findings: Any = item.get("findings") - reads: object = item.get("reads") - if not isinstance(findings, list): - raise RulesError("custom rule host returned invalid findings") - evaluated[(code, _PROJECT_SUBJECT if subject is None else str(subject))] = _Evaluated( - findings=tuple(map(_finding_payload, findings)), - reads=(readsets[reads] if isinstance(reads, int) and code not in untracked else None), - ) - return _HostRun( - evaluated=evaluated, - untracked_reads=untracked, - uncacheable=uncacheable, - observed={ - tuple(map(str, key)): None if digest is None else str(digest) - for key, digest in raw_observed - }, - ) - - def _validate_paths( *, project: CompiledProject, rules: tuple[Rule, ...], - evaluated: dict[tuple[str, str], _Evaluated], + evaluated: dict[tuple[str, str], CustomHostEvaluation], ) -> None: codes: frozenset[str] = frozenset(rule.code for rule in rules) model_rules: frozenset[str] = frozenset( @@ -397,7 +312,7 @@ def _store( stored: dict[str, _RuleCache], identities: dict[str, str], subjects: dict[str, tuple[_Subject, ...]], - run: _HostRun, + run: CustomHostRun, digests: FactDigests, ) -> dict[str, _RuleCache]: updated: dict[str, _RuleCache] = {} @@ -495,7 +410,7 @@ def _read_cache(path: Path) -> dict[str, _RuleCache]: try: stored[str(code)] = _RuleCache( identity=str(value["identity"]), - readsets=list(map(_decode_readset, value["readsets"])), + readsets=list(map(decode_readset, value["readsets"])), entries={ str(subject): _decode_entry(entry) for subject, entry in value["entries"].items() @@ -538,23 +453,15 @@ def _write_cache( temporary.replace(path) -def _decode_readset(value: Any) -> tuple[FactKey, ...]: - return tuple(tuple(map(str, key)) for key in value) - - def _decode_entry(value: Any) -> _CachedEntry: return _CachedEntry( subject_identity=str(value[0]), readset=int(value[1]), reads_digest=str(value[2]), - findings=tuple(map(_finding_payload, value[3])), + findings=tuple(map(finding_payload, value[3])), ) -def _finding_payload(value: Any) -> dict[str, object]: - return {str(key): item for key, item in value.items()} - - def _digest(value: object) -> str: return hashlib.sha256(orjson.dumps(value, default=str)).hexdigest() diff --git a/src/sqlbuild/rule_engine/_helpers/engine/native.py b/src/sqlbuild/rule_engine/_helpers/engine/native.py index a609ba1cc8..dcf33baf81 100644 --- a/src/sqlbuild/rule_engine/_helpers/engine/native.py +++ b/src/sqlbuild/rule_engine/_helpers/engine/native.py @@ -4,7 +4,7 @@ import inspect import json -from concurrent.futures import Future, ThreadPoolExecutor +from concurrent.futures import Executor, Future, ThreadPoolExecutor from dataclasses import asdict, replace from decimal import Decimal from pathlib import Path @@ -64,14 +64,15 @@ def evaluate_native( initial_findings: tuple[Finding, ...] = (), defer_suppressions: bool = False, verify_determinism: bool = False, + custom_payloads: list[dict[str, object]] | None = None, + custom_outcome: Future[CustomRulesOutcome] | None = None, ) -> RulesResult: """Evaluate built-in rules natively while custom rules run incrementally beside them.""" - custom_payloads: list[dict[str, object]] = _custom_rule_payloads( - catalogue=catalogue, project_dir=project_dir - ) - custom_rules: tuple[Rule, ...] = _selected_custom_rules( - config=config, catalogue=catalogue, custom_payloads=custom_payloads + payloads: list[dict[str, object]] = ( + custom_rule_payloads(catalogue=catalogue, project_dir=project_dir) + if custom_payloads is None + else custom_payloads ) request: dict[str, object] = { "version": RULES_NATIVE_API_VERSION, @@ -96,21 +97,22 @@ def evaluate_native( "scope_index": scope_metadata_projection(index=project.scope_index), "initial_findings": [_finding_payload(finding) for finding in initial_findings], "defer_suppressions": True, - "custom_rules": custom_payloads, + "custom_rules": payloads, } with ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-custom-rules") as pool: custom_future: Future[CustomRulesOutcome] | None = ( - pool.submit( - evaluate_custom_rules_cached, + custom_outcome + if custom_payloads is not None + else start_custom_rules( + executor=pool, project=project, config=config, - project_dir=project_dir.resolve(), - rules=custom_rules, + project_dir=project_dir, + catalogue=catalogue, + custom_payloads=payloads, dialect=dialect, verify_determinism=verify_determinism, ) - if custom_rules - else None ) try: response_json: str @@ -157,6 +159,39 @@ def evaluate_native( ) +def start_custom_rules( + *, + executor: Executor, + project: CompiledProject, + config: RulesConfig, + project_dir: Path, + catalogue: tuple[Rule, ...], + custom_payloads: list[dict[str, object]], + dialect: str, + verify_determinism: bool = False, +) -> Future[CustomRulesOutcome] | None: + """Submit selected custom rules, reusing the implementation fingerprints of their payloads.""" + + custom_rules: tuple[Rule, ...] = _selected_custom_rules( + config=config, catalogue=catalogue, custom_payloads=custom_payloads + ) + if not custom_rules: + return None + return executor.submit( + evaluate_custom_rules_cached, + project=project, + config=config, + project_dir=project_dir.resolve(), + rules=custom_rules, + dialect=dialect, + verify_determinism=verify_determinism, + implementation_fingerprints={ + str(payload["code"]): str(payload["implementation_fingerprint"]) + for payload in custom_payloads + }, + ) + + def _selected_custom_rules( *, config: RulesConfig, @@ -261,13 +296,21 @@ def native_catalogue() -> tuple[dict[str, object], ...]: def native_selected_codes( - *, config: RulesConfig, catalogue: tuple[Rule, ...], project_dir: Path + *, + config: RulesConfig, + catalogue: tuple[Rule, ...], + project_dir: Path, + custom_payloads: list[dict[str, object]] | None = None, ) -> tuple[str, ...]: """Resolve the active ruleset through the native Fensu adapter.""" return _selected_codes( config=config, - custom_payloads=_custom_rule_payloads(catalogue=catalogue, project_dir=project_dir), + custom_payloads=( + custom_rule_payloads(catalogue=catalogue, project_dir=project_dir) + if custom_payloads is None + else custom_payloads + ), ) @@ -572,9 +615,11 @@ def _typed_value_payload(value: SqlValue) -> object: } -def _custom_rule_payloads( +def custom_rule_payloads( *, catalogue: tuple[Rule, ...], project_dir: Path ) -> list[dict[str, object]]: + """Describe every catalogued custom rule for native selection and evaluation.""" + closures: dict[str, tuple[Path, ...]] = {} return [ _custom_rule_payload( diff --git a/src/sqlbuild/rule_engine/_helpers/host/custom_evaluation.py b/src/sqlbuild/rule_engine/_helpers/host/custom_evaluation.py index 17fbe73fe7..832e467da7 100644 --- a/src/sqlbuild/rule_engine/_helpers/host/custom_evaluation.py +++ b/src/sqlbuild/rule_engine/_helpers/host/custom_evaluation.py @@ -3,6 +3,7 @@ from __future__ import annotations from collections import Counter +from collections.abc import Callable from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass from itertools import chain @@ -22,12 +23,18 @@ build_rule_fact_views, ) from sqlbuild.rule_engine.classes.runtime_guard import RuntimeGuard +from sqlbuild.rule_engine.constants import ( + CUSTOM_HOST_CANCELLED_MESSAGE, + CUSTOM_RULE_PROJECT_SUBJECT, +) from sqlbuild.rule_engine.exceptions import ( + HostCancelledError, NonHermeticRuleError, OpaqueModuleStateError, RulesError, ) from sqlbuild.rule_engine.models import ( + CustomHostPartition, CustomRuleEvaluation, CustomRuleRun, Finding, @@ -51,6 +58,7 @@ class _PassInputs: facts: RuleFactViews plan: CustomRulePlan | None guard: RuntimeGuard | None + cancelled: Callable[[], bool] | None def evaluate_custom_rules( @@ -64,6 +72,7 @@ def evaluate_custom_rules( track_reads: bool = False, verify_determinism: bool = False, guard: RuntimeGuard | None = None, + partition: CustomHostPartition | None = None, ) -> CustomRuleRun: """Run selected custom rules for every planned subject through the Python authoring API.""" @@ -80,7 +89,10 @@ def evaluate_custom_rules( ) reads: FactReads | None = FactReads() if track_reads else None tracker: _ReadTracker = _ReadTracker( - reads=reads, rules=custom_rules, rules_root=project_dir / "rules" + reads=reads, + rules=custom_rules, + rules_root=project_dir / "rules", + detect_state=partition is not None, ) inputs: _PassInputs = _PassInputs( @@ -92,6 +104,7 @@ def evaluate_custom_rules( facts=facts, plan=plan, guard=guard, + cancelled=None if partition is None else partition.cancelled, ) def run_pass( @@ -118,6 +131,7 @@ def run_pass( untracked_codes=frozenset(tracker.untracked_codes), uncacheable_codes=frozenset(tracker.uncacheable_codes), observed=() if reads is None else tuple(reads.observed.items()), + stateful_codes=frozenset(tracker.stateful_codes), ) @@ -125,15 +139,21 @@ class _ReadTracker: """Attribute reads to one invocation; rules whose module state changes become untracked.""" def __init__( - self, *, reads: FactReads | None, rules: tuple[Rule, ...], rules_root: Path + self, + *, + reads: FactReads | None, + rules: tuple[Rule, ...], + rules_root: Path, + detect_state: bool = False, ) -> None: self.reads: FactReads | None = reads self.untracked_codes: set[str] = set() self.uncacheable_codes: set[str] = set() + self.stateful_codes: set[str] = set() self._rules_root: Path = rules_root self._namespaces: dict[str, tuple[dict[str, object], ...]] = ( {rule.code: rule_namespaces((rule.check,)) for rule in rules} - if reads is not None + if reads is not None or detect_state else {} ) self._states: dict[str, str | None] = { @@ -148,20 +168,22 @@ def start(self) -> None: def finish(self, *, code: str) -> frozenset[tuple[object, ...]] | None: reads: FactReads | None = self.reads - if reads is None: - return None - observed: set[tuple[object, ...]] | None = reads.current - reads.current = None - if reads.untracked: - self.untracked_codes.add(code) - if reads.uncacheable: - self.uncacheable_codes.add(code) + observed: set[tuple[object, ...]] | None = None + if reads is not None: + observed = reads.current + reads.current = None + if reads.untracked: + self.untracked_codes.add(code) + if reads.uncacheable: + self.uncacheable_codes.add(code) for rule_code, namespaces in self._namespaces.items(): - if rule_code in self.untracked_codes: + if rule_code in self.stateful_codes: continue state: str | None = self._state_token(namespaces) if state is None or state != self._states[rule_code]: - self.untracked_codes.add(rule_code) + self.stateful_codes.add(rule_code) + if reads is not None: + self.untracked_codes.add(rule_code) return None if observed is None else frozenset(observed) def _state_token(self, namespaces: tuple[dict[str, object], ...]) -> str | None: @@ -211,7 +233,11 @@ def context(rule: Rule) -> EvaluationRuleContext: (rule, context(rule), None if plan is None else plan.get(rule.code)) for rule in model_rules ) for rule in project_rules: + planned_project: frozenset[str] | None = None if plan is None else plan.get(rule.code) + if planned_project is not None and CUSTOM_RULE_PROJECT_SUBJECT not in planned_project: + continue ctx: EvaluationRuleContext = context(rule) + _stop_if_cancelled(inputs) tracker.start() try: findings: list[Finding] = _invoke( @@ -241,6 +267,7 @@ def context(rule: Rule) -> EvaluationRuleContext: for rule, ctx, planned in model_contexts: if planned is not None and path not in planned: continue + _stop_if_cancelled(inputs) tracker.start() try: findings = _invoke(rule=rule, subject=subject, ctx=ctx, guard=guard) @@ -261,6 +288,13 @@ def context(rule: Rule) -> EvaluationRuleContext: return evaluations +def _stop_if_cancelled(inputs: _PassInputs) -> None: + """Stop between invocations, outside any Rule's guarded scope, once the run is cancelled.""" + + if inputs.cancelled is not None and inputs.cancelled(): + raise HostCancelledError(CUSTOM_HOST_CANCELLED_MESSAGE) + + def _invoke( *, rule: Rule, subject: object, ctx: EvaluationRuleContext, guard: RuntimeGuard | None ) -> list[Finding]: diff --git a/src/sqlbuild/rule_engine/_helpers/host/custom_host.py b/src/sqlbuild/rule_engine/_helpers/host/custom_host.py index ce315b2a0b..f19d2b5347 100644 --- a/src/sqlbuild/rule_engine/_helpers/host/custom_host.py +++ b/src/sqlbuild/rule_engine/_helpers/host/custom_host.py @@ -18,6 +18,8 @@ from sqlbuild.rule_engine._helpers.host.custom_evaluation import evaluate_custom_rules from sqlbuild.rule_engine.classes.runtime_guard import RuntimeGuard from sqlbuild.rule_engine.constants import ( + CUSTOM_HOST_CANCELLED_EXIT_CODE, + CUSTOM_HOST_CANCELLED_MESSAGE, CUSTOM_HOST_HASH_SEED, CUSTOM_HOST_INPUT_TUPLE_SIZE, CUSTOM_HOST_LAUNCH_MODULE, @@ -25,8 +27,14 @@ CUSTOM_HOST_PROTOCOL_VERSION, CUSTOM_HOST_RUNTIME_VERSION, ) -from sqlbuild.rule_engine.exceptions import RulesError -from sqlbuild.rule_engine.models import CustomRuleRun, Finding, Rule, RulesConfig +from sqlbuild.rule_engine.exceptions import HostCancelledError, RulesError +from sqlbuild.rule_engine.models import ( + CustomHostPartition, + CustomRuleRun, + Finding, + Rule, + RulesConfig, +) from sqlbuild.rule_engine.types import CustomRulePlan @@ -50,6 +58,9 @@ def main() -> int: payload: object = request["payload"] if not isinstance(payload, dict): raise RulesError("custom host payload must be an object") + 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_dir: Path = Path(str(payload["project_dir"])).resolve() dialect: str = str(payload.get("dialect", "generic")) @@ -77,7 +88,15 @@ def main() -> int: track_reads=track_reads, verify_determinism=verify_determinism, guard=guard, + partition=( + None + if cancel_marker is None + else CustomHostPartition(cancelled=cancel_marker.exists) + ), ) + except HostCancelledError: + sys.stderr.write(CUSTOM_HOST_CANCELLED_MESSAGE) + return CUSTOM_HOST_CANCELLED_EXIT_CODE except Exception as error: violation: Exception | None = None if guard is None else guard.violation return _write_error(str(violation or error)) @@ -114,6 +133,20 @@ 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.""" + + value: object = payload.get("cancel_marker") + 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") + return marker + + def _decode_plan(value: object) -> CustomRulePlan: if not isinstance(value, dict): raise RulesError("custom host plan must be an object") @@ -145,7 +178,8 @@ def _run_payload(run: CustomRuleRun) -> dict[str, object]: "reads": index, } ) - if sum(map(len, readsets)) > CUSTOM_HOST_MAX_TRACKED_READS: + bounded: bool = sum(map(len, readsets)) > CUSTOM_HOST_MAX_TRACKED_READS + if bounded: untracked.update(str(item["code"]) for item in evaluations if item["reads"] is not None) readsets = {} for item in evaluations: @@ -162,6 +196,8 @@ def _run_payload(run: CustomRuleRun) -> dict[str, object]: }, "uncacheable": sorted(run.uncacheable_codes), "observed": sorted([list(key), digest] for key, digest in run.observed), + "bounded": bounded, + "stateful": sorted(run.stateful_codes), } diff --git a/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py b/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py new file mode 100644 index 0000000000..e84492b29a --- /dev/null +++ b/src/sqlbuild/rule_engine/_helpers/host/custom_host_pool.py @@ -0,0 +1,347 @@ +"""Run planned custom-rule invocations across isolated host processes and merge their results.""" + +from __future__ import annotations + +import pickle +import shutil +import sys +import tempfile +from concurrent.futures import FIRST_EXCEPTION, Future, ThreadPoolExecutor, wait +from dataclasses import replace +from pathlib import Path +from typing import Any, cast + +import orjson + +import sqlbuild._native as _native +from sqlbuild.compiler.compile.models import CompiledProject +from sqlbuild.rule_engine._helpers.host.host_capacity import available_cores +from sqlbuild.rule_engine.constants import ( + CUSTOM_HOST_CANCELLED_MARKER, + CUSTOM_HOST_LAUNCH_MODULE, + CUSTOM_HOST_MAX_TRACKED_READS, + CUSTOM_HOST_RUNTIME_VERSION, + CUSTOM_RULE_PROJECT_SUBJECT, +) +from sqlbuild.rule_engine.exceptions import RulesError +from sqlbuild.rule_engine.models import ( + CustomHostEvaluation, + CustomHostRun, + CustomHostSlice, + Rule, + RulesConfig, +) +from sqlbuild.rule_engine.types import CustomHostPlan, FactKey, RuleSubject + +_HOST_TIMEOUT_MILLIS: int = 120_000 +_MIN_INVOCATIONS_PER_HOST: int = 64 +_MAX_HOSTS: int = 8 + + +def run_custom_hosts( + *, + project: CompiledProject, + config: RulesConfig, + project_dir: Path, + dialect: str, + rules: tuple[Rule, ...], + model_paths: tuple[str, ...], + plan: CustomHostPlan, + track_reads: bool, + verify_determinism: bool, +) -> CustomHostRun: + """Evaluate the plan across host processes and merge their results in single-host order.""" + + invocations: tuple[tuple[str, str], ...] = planned_invocations( + plan=plan, rules=rules, model_paths=model_paths + ) + plans: tuple[CustomHostPlan, ...] = ( + (plan,) + if verify_determinism + else partition_host_plan( + plan=plan, invocations=invocations, hosts=host_count(len(invocations)) + ) + ) + input_dir: Path = project_dir / "target" / "rules-cache" / "host-inputs" + input_dir.mkdir(parents=True, exist_ok=True) + input_path: Path | None = None + 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) + ) + finally: + if input_path is not None: + input_path.unlink(missing_ok=True) + return _merge( + responses=responses, + order={invocation: index for index, invocation in enumerate(invocations)}, + ) + + +def _run_partitions( + *, + payload: dict[str, object], + plan: CustomHostPlan, + plans: tuple[CustomHostPlan, ...], + input_dir: Path, +) -> list[CustomHostSlice]: + """Run partitions; failures and rules with module state fall back to one host.""" + + cancel_dir: Path = Path(tempfile.mkdtemp(dir=input_dir, prefix="hosts-")) + marker: Path = cancel_dir / CUSTOM_HOST_CANCELLED_MARKER + futures: list[Future[CustomHostSlice]] = [] + failed: bool = False + try: + with ThreadPoolExecutor( + max_workers=len(plans), thread_name_prefix="sqlbuild-custom-host" + ) as pool: + futures = [ + pool.submit( + _run_host, + payload={ + **payload, + "plan": partition, + "cancel_marker": str(marker.resolve()), + }, + ) + for partition in plans + ] + done: set[Future[CustomHostSlice]] = wait(futures, return_when=FIRST_EXCEPTION).done + failed = any(future.exception() is not None for future in done) + if failed: + _cancel_hosts(marker) + finally: + shutil.rmtree(cancel_dir, ignore_errors=True) + if failed: + return [_run_host(payload={**payload, "plan": plan})] + slices: list[CustomHostSlice] = [future.result() for future in futures] + stateful: frozenset[str] = frozenset().union(*(item.stateful for item in slices)) + if not stateful: + return slices + single: 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] + + +def _cancel_hosts(marker: Path) -> None: + """Stop hosts at their next invocation; a running invocation still ends or times out.""" + + marker.touch() + + +def _without_codes(*, response: CustomHostSlice, codes: frozenset[str]) -> CustomHostSlice: + """Drop every result a host reported for the given rules.""" + + return replace( + response, + evaluations=tuple(item for item in response.evaluations if item[0] not in codes), + untracked_reads={ + code: reads for code, reads in response.untracked_reads.items() if code not in codes + }, + uncacheable=response.uncacheable - codes, + stateful=response.stateful - codes, + ) + + +def planned_invocations( + *, plan: CustomHostPlan, rules: tuple[Rule, ...], model_paths: tuple[str, ...] +) -> tuple[tuple[str, str], ...]: + """List planned (rule, subject) invocations in the order one host evaluates them.""" + + planned: tuple[Rule, ...] = tuple( + sorted((rule for rule in rules if rule.code in plan), key=lambda rule: rule.code) + ) + invocations: list[tuple[str, str]] = [ + (rule.code, CUSTOM_RULE_PROJECT_SUBJECT) + for rule in planned + if rule.subject is RuleSubject.PROJECT + ] + model_rules: tuple[tuple[str, frozenset[str] | None], ...] = tuple( + (rule.code, None if plan[rule.code] is None else frozenset(plan[rule.code] or ())) + for rule in planned + if rule.subject is RuleSubject.MODEL + ) + for path in model_paths if model_rules else (): + invocations.extend( + (code, path) for code, selected in model_rules if selected is None or path in selected + ) + return tuple(invocations) + + +def partition_host_plan( + *, plan: CustomHostPlan, invocations: tuple[tuple[str, str], ...], hosts: int +) -> tuple[CustomHostPlan, ...]: + """Deal project invocations and whole model subjects round-robin across at most `hosts`.""" + + units: list[list[tuple[str, str]]] = [] + for code, subject in invocations: + if subject == CUSTOM_RULE_PROJECT_SUBJECT or not units or units[-1][0][1] != subject: + units.append([]) + units[-1].append((code, subject)) + count: int = max(1, min(hosts, len(units))) + if count == 1: + return (plan,) + partitions: list[dict[str, list[str]]] = [] + for _ in range(count): + partitions.append({code: [] for code in plan}) + for index, unit in enumerate(units): + for code, subject in unit: + partitions[index % count][code].append(subject) + return tuple(cast(CustomHostPlan, partition) for partition in partitions) + + +def host_count(invocations: int) -> int: + """Hosts worth starting: one per 64 invocations, bounded by usable CPUs and eight.""" + + return max(1, min(_MAX_HOSTS, available_cores(), invocations // _MIN_INVOCATIONS_PER_HOST)) + + +def _run_host(*, payload: dict[str, object]) -> CustomHostSlice: + try: + response_json: str = _native.run_custom_host_json( + orjson.dumps( + { + "program": sys.executable, + "arguments": ["-m", CUSTOM_HOST_LAUNCH_MODULE], + "timeout_millis": _HOST_TIMEOUT_MILLIS, + "runtime_version": CUSTOM_HOST_RUNTIME_VERSION, + "payload": payload, + } + ).decode() + ) + except (ValueError, TypeError) as error: + raise RulesError(str(error)) from error + return _decode_host_response(orjson.loads(response_json)) + + +def _decode_host_response(response: object) -> CustomHostSlice: + if not isinstance(response, dict): + raise RulesError("custom rule host returned an invalid payload") + payload: dict[str, Any] = cast(dict[str, Any], response) + raw_readsets: Any = payload.get("readsets") + raw_evaluations: Any = payload.get("evaluations") + raw_untracked: Any = payload.get("untracked") + raw_uncacheable: Any = payload.get("uncacheable") + raw_observed: Any = payload.get("observed") + if ( + not isinstance(raw_readsets, list) + or not isinstance(raw_evaluations, list) + or not isinstance(raw_untracked, dict) + or not isinstance(raw_uncacheable, list) + or not isinstance(raw_observed, list) + ): + raise RulesError("custom rule host returned an invalid payload") + readsets: list[tuple[FactKey, ...]] = list(map(decode_readset, raw_readsets)) + untracked: dict[str, tuple[FactKey, ...]] = { + str(code): decode_readset(reads) for code, reads in raw_untracked.items() + } + evaluations: list[tuple[str, str, CustomHostEvaluation]] = [] + for raw_item in raw_evaluations: + if not isinstance(raw_item, dict): + raise RulesError("custom rule host returned an invalid evaluation") + item: dict[str, Any] = cast(dict[str, Any], raw_item) + code: str = str(item["code"]) + subject: object = item.get("subject") + findings: Any = item.get("findings") + reads: object = item.get("reads") + if not isinstance(findings, list): + raise RulesError("custom rule host returned invalid findings") + tracked: bool = isinstance(reads, int) and code not in untracked + evaluations.append( + ( + code, + CUSTOM_RULE_PROJECT_SUBJECT if subject is None else str(subject), + CustomHostEvaluation( + findings=tuple(map(finding_payload, findings)), + reads=readsets[cast(int, reads)] if tracked else None, + ), + ) + ) + return CustomHostSlice( + evaluations=tuple(evaluations), + untracked_reads=untracked, + uncacheable=frozenset(str(code) for code in raw_uncacheable), + observed={ + tuple(map(str, key)): None if digest is None else str(digest) + for key, digest in raw_observed + }, + bounded=payload.get("bounded") is True, + stateful=frozenset(str(code) for code in payload.get("stateful") or ()), + ) + + +def _merge(*, responses: list[CustomHostSlice], order: dict[tuple[str, str], int]) -> CustomHostRun: + """Combine host slices exactly as one host would have reported every invocation.""" + + evaluations: list[tuple[str, str, CustomHostEvaluation]] = [] + untracked: set[str] = set() + for response in responses: + evaluations.extend(response.evaluations) + untracked.update(response.untracked_reads) + if len(responses) > 1: + evaluations.sort(key=lambda item: order.get((item[0], item[1]), len(order))) + tracked_readsets: set[tuple[FactKey, ...]] = { + result.reads + for code, _subject, result in evaluations + if result.reads is not None and code not in untracked + } + if any(response.bounded for response in responses) or ( + sum(map(len, tracked_readsets)) > CUSTOM_HOST_MAX_TRACKED_READS + ): + untracked.update(code for code, _subject, result in evaluations if result.reads is not None) + untracked_reads: dict[str, set[FactKey]] = {code: set() for code in untracked} + for response in responses: + for code, reads in response.untracked_reads.items(): + untracked_reads[code].update(reads) + evaluated: dict[tuple[str, str], CustomHostEvaluation] = {} + for code, subject, result in evaluations: + merged: CustomHostEvaluation = result + if code in untracked and result.reads is not None: + untracked_reads[code].update(result.reads) + merged = CustomHostEvaluation(findings=result.findings, reads=None) + evaluated[(code, subject)] = merged + return CustomHostRun( + evaluated=evaluated, + untracked_reads={code: tuple(sorted(reads)) for code, reads in untracked_reads.items()}, + uncacheable=frozenset().union(*(response.uncacheable for response in responses)), + observed=_merged_observations(responses), + ) + + +def _merged_observations(responses: list[CustomHostSlice]) -> dict[FactKey, str | None]: + """Keep each observed digest, forgetting any that two hosts observed differently.""" + + observed: dict[FactKey, str | None] = {} + for response in responses: + for key, digest in response.observed.items(): + if key in observed and observed[key] != digest: + observed[key] = None + else: + observed.setdefault(key, digest) + return observed + + +def decode_readset(value: Any) -> tuple[FactKey, ...]: + return tuple(tuple(map(str, key)) for key in value) + + +def finding_payload(value: Any) -> dict[str, object]: + return {str(key): item for key, item in value.items()} diff --git a/src/sqlbuild/rule_engine/_helpers/host/host_capacity.py b/src/sqlbuild/rule_engine/_helpers/host/host_capacity.py new file mode 100644 index 0000000000..0f6d4e5840 --- /dev/null +++ b/src/sqlbuild/rule_engine/_helpers/host/host_capacity.py @@ -0,0 +1,103 @@ +"""CPU capacity available to this process, including container CPU quotas.""" + +from __future__ import annotations + +import math +import os +from collections.abc import Callable, Iterator +from pathlib import Path + +_CGROUP_ROOT: Path = Path("/sys/fs/cgroup") +_PROC_CGROUP: Path = Path("/proc/self/cgroup") +_CGROUP_V1_CPU_CONTROLLER: str = "cpu" +_CGROUP_V1_CPU_DIRECTORIES: tuple[str, ...] = ("cpu", "cpu,cpuacct", "cpuacct,cpu") +_CPU_MAX_FIELD_COUNT: int = 2 + + +def available_cores(*, cgroup_root: Path = _CGROUP_ROOT, proc_cgroup: Path = _PROC_CGROUP) -> int: + """Return usable CPUs: the process CPU count capped by any cgroup v1 or v2 CPU quota.""" + + counter: Callable[[], int | None] | None = getattr(os, "process_cpu_count", None) + count: int = (counter() if counter is not None else _affinity_count()) or 1 + quotas: list[int] = list(_cgroup_quotas(cgroup_root=cgroup_root, proc_cgroup=proc_cgroup)) + return max(1, min((count, *quotas))) + + +def _affinity_count() -> int | None: + affinity: Callable[[int], set[int]] | None = getattr(os, "sched_getaffinity", None) + return len(affinity(0)) if affinity is not None else os.cpu_count() + + +def _cgroup_quotas(*, cgroup_root: Path, proc_cgroup: Path) -> Iterator[int]: + try: + entries: list[str] = proc_cgroup.read_text(encoding="utf-8").splitlines() + except (OSError, ValueError): + entries = [] + for entry in entries or ["0::/"]: + _hierarchy, _separator, remainder = entry.partition(":") + controllers, _separator, path = remainder.partition(":") + if not controllers: + yield from _quotas_in( + paths=tuple( + directory / "cpu.max" + for directory in _directories(base=cgroup_root, relative=path) + ), + parse=_cpu_max_quota, + ) + elif _CGROUP_V1_CPU_CONTROLLER in controllers.split(","): + for name in (controllers, *_CGROUP_V1_CPU_DIRECTORIES): + yield from _quotas_in( + paths=tuple(_directories(base=cgroup_root / name, relative=path)), + parse=_cfs_quota, + ) + + +def _directories(*, base: Path, relative: str) -> Iterator[Path]: + """Yield the cgroup directory and each ancestor down from the hierarchy root.""" + + directory: Path = base + yield directory + for part in Path(relative.strip("/")).parts: + directory = directory / part + yield directory + + +def _quotas_in(*, paths: tuple[Path, ...], parse: Callable[[Path], int | None]) -> Iterator[int]: + for path in paths: + quota: int | None = parse(path) + if quota is not None: + yield quota + + +def _cpu_max_quota(path: Path) -> int | None: + """Parse cgroup v2 `cpu.max`: `max ` means unlimited, else ` `.""" + + try: + fields: list[str] = path.read_text(encoding="utf-8").split() + except (OSError, ValueError): + return None + return ( + _quota(quota=fields[0], period=fields[1]) if len(fields) == _CPU_MAX_FIELD_COUNT else None + ) + + +def _cfs_quota(directory: Path) -> int | None: + """Parse cgroup v1 `cpu.cfs_quota_us` / `cpu.cfs_period_us`; a negative quota is unlimited.""" + + try: + quota: str = (directory / "cpu.cfs_quota_us").read_text(encoding="utf-8").strip() + period: str = (directory / "cpu.cfs_period_us").read_text(encoding="utf-8").strip() + except (OSError, ValueError): + return None + return _quota(quota=quota, period=period) + + +def _quota(*, quota: str, period: str) -> int | None: + try: + quota_us: int = int(quota) + period_us: int = int(period) + except ValueError: + return None + if quota_us <= 0 or period_us <= 0: + return None + return max(1, math.ceil(quota_us / period_us)) diff --git a/src/sqlbuild/rule_engine/_helpers/run/rules.py b/src/sqlbuild/rule_engine/_helpers/run/rules.py index a9aa96dd97..a4053439f7 100644 --- a/src/sqlbuild/rule_engine/_helpers/run/rules.py +++ b/src/sqlbuild/rule_engine/_helpers/run/rules.py @@ -51,10 +51,12 @@ from sqlbuild.rule_engine._helpers.engine.config import resolve_rule_ignore_selectors from sqlbuild.rule_engine._helpers.engine.hermeticity import verify_custom_rules from sqlbuild.rule_engine._helpers.engine.native import ( + custom_rule_payloads, decode_rule_finding, evaluate_native, finalize_native_findings, native_catalogue, + start_custom_rules, ) from sqlbuild.rule_engine._helpers.run.findings import group_unevaluated_findings from sqlbuild.rule_engine._helpers.run.literal_duplicates import with_duplicate_literal_hints @@ -62,6 +64,7 @@ from sqlbuild.rule_engine.exceptions import RulesError from sqlbuild.rule_engine.main.load_config import load_rules_config from sqlbuild.rule_engine.models import ( + CustomRulesOutcome, Finding, PreparedSqlLint, PreparedSqlLintResult, @@ -114,8 +117,14 @@ def evaluate_rules( project_dir=resolved_project_dir, include_custom=include_custom, ) + custom_payloads: list[dict[str, object]] = custom_rule_payloads( + catalogue=catalogue, project_dir=resolved_project_dir + ) selected: tuple[Rule, ...] = select_rules( - catalogue=catalogue, config=effective_config, project_dir=resolved_project_dir + catalogue=catalogue, + config=effective_config, + project_dir=resolved_project_dir, + custom_payloads=custom_payloads, ) skipped_type_proof_rules: tuple[str, ...] = () if no_sql_analysis or not graph.project.settings.sql_analysis: @@ -140,20 +149,19 @@ def evaluate_rules( select=tuple(rule.code for rule in selected), ignore=(), ) - lint_config: LintConfig = sql_lint_config( - dialect=dialect, - codes=_sql_rule_codes(native_rules), - thresholds=effective_config.thresholds, - relation_keys=compiled_relation_keys(graph.project), - ) - sql_rule_inputs: _SqlRuleInputs = _prepare_sql_rule_inputs( - rules=native_rules, - project=selected_project, - dialect=_lint_identity(lint_config), - project_dir=resolved_project_dir, - collect_files=prepared_sql is None and model_paths is None, - ) - with ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-rules") as executor: + with ( + ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-custom-rules") as custom, + ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-rules") as executor, + ): + custom_outcome: Future[CustomRulesOutcome] | None = start_custom_rules( + executor=custom, + project=selected_project, + config=selected_config, + project_dir=resolved_project_dir, + catalogue=catalogue, + custom_payloads=custom_payloads, + dialect=dialect, + ) native_result: Future[RulesResult] = executor.submit( evaluate_native, project=selected_project, @@ -162,6 +170,21 @@ def evaluate_rules( catalogue=catalogue, dialect=dialect, defer_suppressions=True, + custom_payloads=custom_payloads, + custom_outcome=custom_outcome, + ) + lint_config: LintConfig = sql_lint_config( + dialect=dialect, + codes=_sql_rule_codes(native_rules), + thresholds=effective_config.thresholds, + relation_keys=compiled_relation_keys(graph.project), + ) + sql_rule_inputs: _SqlRuleInputs = _prepare_sql_rule_inputs( + rules=native_rules, + project=selected_project, + dialect=_lint_identity(lint_config), + project_dir=resolved_project_dir, + collect_files=prepared_sql is None and model_paths is None, ) sql_started: float = time.monotonic() sql_result: _SqlRulesEvaluation = _run_sql_rules( @@ -217,7 +240,7 @@ def _prepare_sql_rule_inputs( project_dir: Path, collect_files: bool, ) -> _SqlRuleInputs: - """Hash SQL and read files first so GIL-releasing work never waits on the native thread.""" + """Hash SQL and read the project files on the caller while built-in rules evaluate natively.""" identities: dict[str, str] = _sql_model_rule_identities( rules=rules, project=project, dialect=dialect diff --git a/src/sqlbuild/rule_engine/constants.py b/src/sqlbuild/rule_engine/constants.py index b9e266f999..98078ffac1 100644 --- a/src/sqlbuild/rule_engine/constants.py +++ b/src/sqlbuild/rule_engine/constants.py @@ -10,12 +10,16 @@ CUSTOM_HOST_PROTOCOL_VERSION: int = 1 CUSTOM_HOST_RUNTIME_VERSION: str = "sqlbuild-rules-custom-v2" CUSTOM_HOST_MAX_TRACKED_READS: int = 100_000 +CUSTOM_RULE_PROJECT_SUBJECT: str = "" CUSTOM_RULES_CACHE_VERSION: str = "custom-rules-reads-v1" CUSTOM_RULES_CACHE_FILE: str = "target/rules-cache/bulk/custom-rules.json" NATIVE_RULES_MEMO_VERSION: str = "native-rules-response-v1" NATIVE_RULES_MEMO_FILE: str = "target/rules-cache/bulk/native-response.json" CUSTOM_HOST_INPUT_TUPLE_SIZE: int = 2 CUSTOM_HOST_MODULE: str = "sqlbuild.rule_engine._helpers.host.custom_host" +CUSTOM_HOST_CANCELLED_MARKER: str = "cancelled" +CUSTOM_HOST_CANCELLED_EXIT_CODE: int = 75 +CUSTOM_HOST_CANCELLED_MESSAGE: str = "custom host cancelled after another host failed" CUSTOM_HOST_LAUNCH_MODULE: str = "sqlbuild.rule_engine._helpers.host.custom_host_launch" CUSTOM_HOST_HASH_SEED: str = "0" CUSTOM_HOST_EXEC_OS_NAME: str = "posix" diff --git a/src/sqlbuild/rule_engine/exceptions.py b/src/sqlbuild/rule_engine/exceptions.py index 2d164a99a2..9786c0c1f5 100644 --- a/src/sqlbuild/rule_engine/exceptions.py +++ b/src/sqlbuild/rule_engine/exceptions.py @@ -23,3 +23,7 @@ class FactDigestError(RulesError): class OpaqueModuleStateError(RulesError): """Raised when custom-rule module state holds a value whose mutations cannot be observed.""" + + +class HostCancelledError(RulesError): + """Raised in a custom-rule host that stops because another host already failed.""" diff --git a/src/sqlbuild/rule_engine/models.py b/src/sqlbuild/rule_engine/models.py index cb745b884b..99055f59b4 100644 --- a/src/sqlbuild/rule_engine/models.py +++ b/src/sqlbuild/rule_engine/models.py @@ -2,6 +2,7 @@ from __future__ import annotations +from collections.abc import Callable from concurrent.futures import Future from dataclasses import dataclass, field from pathlib import Path @@ -10,7 +11,7 @@ from sqlbuild.compiler.compile.models import DeclarationScopeBuild, SqlExpansionContext from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs from sqlbuild.lint.models import LintRunResult -from sqlbuild.rule_engine.types import RuleCheck, RuleOptionValue, RuleSubject +from sqlbuild.rule_engine.types import FactKey, RuleCheck, RuleOptionValue, RuleSubject @dataclass(frozen=True, slots=True) @@ -401,6 +402,7 @@ class CustomRuleRun: untracked_codes: frozenset[str] uncacheable_codes: frozenset[str] = frozenset() observed: tuple[tuple[tuple[str, ...], str | None], ...] = () + stateful_codes: frozenset[str] = frozenset() @dataclass(frozen=True) @@ -411,3 +413,40 @@ class CustomRulesOutcome: cache_hits: int cache_misses: int custom_ms: int + + +@dataclass(frozen=True) +class CustomHostEvaluation: + """Decoded findings and recorded reads of one custom-rule invocation.""" + + findings: tuple[dict[str, object], ...] + reads: tuple[FactKey, ...] | None + + +@dataclass(frozen=True) +class CustomHostRun: + """Every planned custom-rule invocation, merged across host processes.""" + + evaluated: dict[tuple[str, str], CustomHostEvaluation] + untracked_reads: dict[str, tuple[FactKey, ...]] + uncacheable: frozenset[str] + observed: dict[FactKey, str | None] + + +@dataclass(frozen=True) +class CustomHostSlice: + """One host process's invocations, in the order that host evaluated them.""" + + evaluations: tuple[tuple[str, str, CustomHostEvaluation], ...] + untracked_reads: dict[str, tuple[FactKey, ...]] + uncacheable: frozenset[str] + observed: dict[FactKey, str | None] + bounded: bool + stateful: frozenset[str] + + +@dataclass(frozen=True) +class CustomHostPartition: + """How one host of a split run detects module state and learns the run was cancelled.""" + + cancelled: Callable[[], bool] diff --git a/src/sqlbuild/rule_engine/types.py b/src/sqlbuild/rule_engine/types.py index ccaedcbba2..477c7e123c 100644 --- a/src/sqlbuild/rule_engine/types.py +++ b/src/sqlbuild/rule_engine/types.py @@ -28,6 +28,8 @@ type FactKey = tuple[str, ...] type CustomRulePlan = dict[str, frozenset[str] | None] +type CustomHostPlan = dict[str, list[str] | None] + class RuleSubject(StrEnum): """Compiler-owned unit used to invoke and cache one custom rule.""" diff --git a/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py b/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py index 71b10cd215..af6a8c559d 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py @@ -156,6 +156,16 @@ class RulesIntegrationTestCase: expected_code: str +@dataclass(frozen=True) +class CustomHostSplitIntegrationTestCase: + """A compile whose custom rules run on several hosts and must match one host.""" + + description: str + model_count: int + hosts: int + expected_code: str + + @dataclass(frozen=True) class DollarQuotedLiteralBuildTestCase: """One model projection with dollar-quoted literals and its built rows.""" diff --git a/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py b/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py index 0ec88e2af8..c8c4cb3bc5 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/test_rules.py @@ -11,10 +11,11 @@ import pytest from _pytest.capture import CaptureResult +import sqlbuild._native as native_module import sqlbuild.compiler.compile._helpers.attachment.declaration_scope as declaration_scope_module import sqlbuild.compiler.compile.main._build_compile_inputs as compile_inputs_module import sqlbuild.compiler.compile.main.sql_expansion_context as expansion_context_module -import sqlbuild.rule_engine._helpers.engine.custom_rules as custom_rules_module +import sqlbuild.rule_engine._helpers.host.custom_host_pool as custom_host_pool_module import sqlbuild.rule_engine._helpers.run.rules as rules_module from sqlbuild.cli.commands.main.entrypoint.entry import main from sqlbuild.compiler.compile.models import DeclarationScopeBuild, LoadedMacro @@ -23,6 +24,7 @@ from sqlbuild.lint.classes.stop_context import LintStopContext from sqlbuild.lint.models import LintRunResult from tests.integration.src.sqlbuild.cli.commands.main._test_types import ( + CustomHostSplitIntegrationTestCase, DynamicPivotRulesIntegrationTestCase, ExplicitContractOutputRuleIntegrationTestCase, ImplicitAliasRuleIntegrationTestCase, @@ -1018,13 +1020,13 @@ def final_directory(*, model: Model, ctx: RuleContext) -> list[Finding]: encoding="utf-8", ) host_runs: list[str] = [] - original_host: Callable[[str], str] = custom_rules_module._native.run_custom_host_json + original_host: Callable[[str], str] = native_module.run_custom_host_json def recording_host(spec_json: str) -> str: host_runs.append(spec_json) return original_host(spec_json) - monkeypatch.setattr(custom_rules_module._native, "run_custom_host_json", recording_host) + monkeypatch.setattr(native_module, "run_custom_host_json", recording_host) cold: set[str] = compile_finding_keys(project_dir=tmp_path, capsys=capsys) @@ -1038,6 +1040,78 @@ def recording_host(spec_json: str) -> str: } <= cold +@pytest.mark.parametrize( + "test_case", + [CustomHostSplitIntegrationTestCase("custom rules split across hosts", 9, 3, "XSQBRARCH001")], + ids=lambda case: case.description, +) +def test_given_custom_rules_split_across_hosts_when_compiling_then_output_matches_one_host( + test_case: CustomHostSplitIntegrationTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], + monkeypatch: pytest.MonkeyPatch, +) -> None: + (tmp_path / "sqlbuild_project.toml").write_text( + 'name = "orders"\nadapter = "duckdb"\n\n' + f'[rules]\nselect = ["{test_case.expected_code}"]\n\n' + "[rules.thresholds]\nmin_custom_rule_test_cases = 0\n", + encoding="utf-8", + ) + for index in range(test_case.model_count): + folder: str = ("final", "staging", "staging")[index % 3] + model: Path = tmp_path / "models" / folder / f"orders_{index:02d}.sql" + model.parent.mkdir(parents=True, exist_ok=True) + model.write_text( + f'MODEL (description "Orders {index}");\nSELECT {index} AS order_id\n', + encoding="utf-8", + ) + rule_file: Path = tmp_path / "rules" / "architecture.py" + rule_file.parent.mkdir() + rule_file.write_text( + """from sqlbuild.rules import Finding, Model, RuleContext, rule + +@rule( + code="XSQBRARCH001", + message="Final models must use the final directory", + remediation="Move this model beneath models/final/.", +) +def final_directory(*, model: Model, ctx: RuleContext) -> list[Finding]: + return [] if "final" in model.path.parts else [ctx.finding(subject=model)] +""", + encoding="utf-8", + ) + host_runs: list[str] = [] + original_host: Callable[[str], str] = native_module.run_custom_host_json + + def recording_host(spec_json: str) -> str: + host_runs.append(spec_json) + return original_host(spec_json) + + def compile_payload(*, hosts: int) -> dict[str, Any]: + monkeypatch.setattr(custom_host_pool_module, "_MIN_INVOCATIONS_PER_HOST", 1) + monkeypatch.setattr(custom_host_pool_module, "available_cores", lambda: hosts) + exit_code: int = main(["--project-dir", str(tmp_path), "compile", "--json", "--no-cache"]) + assert exit_code == 1 + return json.loads(capsys.readouterr().out) + + monkeypatch.setattr(native_module, "run_custom_host_json", recording_host) + single: dict[str, Any] = compile_payload(hosts=1) + single_runs: int = len(host_runs) + (tmp_path / "target" / "rules-cache").rename(tmp_path / "single-rules-cache") + split: dict[str, Any] = compile_payload(hosts=test_case.hosts) + split_runs: int = len(host_runs) - single_runs + warm: dict[str, Any] = compile_payload(hosts=test_case.hosts) + + assert (single_runs, split_runs, len(host_runs)) == (1, test_case.hosts, 1 + test_case.hosts) + assert split["diagnostics"] == single["diagnostics"] + assert len(split["diagnostics"]) == test_case.model_count - test_case.model_count // 3 + assert warm["diagnostics"] == single["diagnostics"] + assert warm["compile_timings"]["rule_cache_hits"] >= test_case.model_count + assert (tmp_path / "single-rules-cache" / "bulk" / "custom-rules.json").read_bytes() == ( + tmp_path / "target" / "rules-cache" / "bulk" / "custom-rules.json" + ).read_bytes() + + @pytest.mark.parametrize( "test_case", [RulesIntegrationTestCase("custom rule is enforced and cached", 1, "XSQBRARCH001")], diff --git a/tests/unit/src/sqlbuild/lint/_helpers/_test_types.py b/tests/unit/src/sqlbuild/lint/_helpers/_test_types.py index 3b4427709d..b49d2a90bb 100644 --- a/tests/unit/src/sqlbuild/lint/_helpers/_test_types.py +++ b/tests/unit/src/sqlbuild/lint/_helpers/_test_types.py @@ -215,3 +215,22 @@ class ManyViolationLocationTestCase: unused_cte_count: int file_count: int expected_code: str + + +@dataclass(frozen=True) +class BatchedPreparationTestCase: + """A project whose lint is compared between batched and per-body native preparation.""" + + description: str + files: dict[str, str] + expected_codes: tuple[str, ...] + + +@dataclass(frozen=True) +class BatchedPreparationFailureTestCase: + """A native preparation failure and an unexpandable file, in a given file order.""" + + description: str + unexpandable_file: str + skip_unexpandable: bool + expected_error: type[Exception] | None diff --git a/tests/unit/src/sqlbuild/lint/_helpers/helpers.py b/tests/unit/src/sqlbuild/lint/_helpers/helpers.py index 9863d1f4fb..cf3135758f 100644 --- a/tests/unit/src/sqlbuild/lint/_helpers/helpers.py +++ b/tests/unit/src/sqlbuild/lint/_helpers/helpers.py @@ -6,6 +6,7 @@ from sqlbuild.lint._helpers.headers import scan_headers, sql_body_ranges from sqlbuild.lint.models import HeaderSpan, LintBody +from sqlbuild.lint.types import NativeLintPreparationRequest, NativePreparedSql def lint_bodies_for(*, file_path: Path, contents: str) -> tuple[LintBody, ...]: @@ -64,3 +65,29 @@ def authored_positions(*, contents: str, needles: tuple[str, ...]) -> list[tuple (contents.count("\n", 0, offset) + 1, offset - contents.rfind("\n", 0, offset)) ) return positions + + +def write_lint_project(*, root: Path, files: dict[str, str]) -> None: + """Write a DuckDB project containing the given relative files.""" + + (root / "sqlbuild_project.toml").write_text( + 'name = "orders"\nadapter = "duckdb"\n', encoding="utf-8" + ) + 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") + + +def prepare_nothing_in_batch( + requests: list[NativeLintPreparationRequest], +) -> list[tuple[bool, NativePreparedSql | None]]: + """Report every body as unprepared so each one takes the per-body native call.""" + + return [(False, None) for _request in requests] + + +def refuse_preparation(request: NativeLintPreparationRequest) -> NativePreparedSql | None: + """Fail the per-body native preparation of any body.""" + + raise ValueError(f"cannot prepare {request['dialect']} body") diff --git a/tests/unit/src/sqlbuild/lint/_helpers/test_expansion.py b/tests/unit/src/sqlbuild/lint/_helpers/test_expansion.py new file mode 100644 index 0000000000..61b810260e --- /dev/null +++ b/tests/unit/src/sqlbuild/lint/_helpers/test_expansion.py @@ -0,0 +1,116 @@ +"""Lint bodies are prepared natively in one batch with per-body results and failure order.""" + +from __future__ import annotations + +from pathlib import Path + +import pytest + +import sqlbuild._native as native_module +from sqlbuild.lint.exceptions import ProjectCompileError +from sqlbuild.lint.main.run_lint import run_lint +from sqlbuild.lint.models import LintConfig, LintRunResult +from tests.unit.src.sqlbuild.lint._helpers._test_types import ( + BatchedPreparationFailureTestCase, + BatchedPreparationTestCase, +) +from tests.unit.src.sqlbuild.lint._helpers.helpers import ( + prepare_nothing_in_batch, + refuse_preparation, + write_lint_project, +) + +_ORDERS_MODEL: str = ( + 'MODEL (description "Orders");\n' + 'SELECT o.customer_id FROM __ref("customers") AS o WHERE o.region = NULL\n' +) +_CUSTOMERS_MODEL: str = ( + "MODEL (description \"Customers\");\nSELECT 1 AS customer_id, 'west' AS region\n" +) +_ORDERS_TEST: str = ( + "TEST ();\nWITH\n__ref__customers AS (SELECT 1 AS customer_id, 'west' AS region),\n" + "__expected__orders AS (SELECT 1 AS customer_id)\nSELECT 1\n" +) +_NOT_NULL_AUDIT: str = "AUDIT ();\nSELECT 1 FROM @relation WHERE @column IS NULL\n" + + +@pytest.mark.parametrize( + "test_case", + [ + BatchedPreparationTestCase( + description="models, tests, and generic audits", + files={ + "models/customers.sql": _CUSTOMERS_MODEL, + "models/orders.sql": _ORDERS_MODEL, + "tests/unit/test_orders.sql": _ORDERS_TEST, + "audits/generic/not_null.sql": _NOT_NULL_AUDIT, + }, + expected_codes=("SQBRSQL001",), + ) + ], + ids=lambda case: case.description, +) +def test_given_project_when_linting_with_batched_preparation_then_result_matches_per_body( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: BatchedPreparationTestCase +) -> None: + write_lint_project(root=tmp_path, files=test_case.files) + config: LintConfig = LintConfig(dialect="duckdb") + + batched: LintRunResult = run_lint(project_dir=tmp_path, config=config) + monkeypatch.setattr(native_module, "prepare_lint_sql_batch", prepare_nothing_in_batch) + per_body: LintRunResult = run_lint(project_dir=tmp_path, config=config) + + assert batched == per_body + assert tuple(violation.code for violation in batched.violations) == test_case.expected_codes + + +@pytest.mark.parametrize( + "test_case", + [ + BatchedPreparationFailureTestCase( + description="native failure before a later unexpandable file", + unexpandable_file="models/zz_broken.sql", + skip_unexpandable=False, + expected_error=ValueError, + ), + BatchedPreparationFailureTestCase( + description="unexpandable file before a later native failure", + unexpandable_file="models/aa_broken.sql", + skip_unexpandable=False, + expected_error=ProjectCompileError, + ), + BatchedPreparationFailureTestCase( + description="skipped unexpandable file before a later native failure", + unexpandable_file="models/aa_broken.sql", + skip_unexpandable=True, + expected_error=ValueError, + ), + ], + ids=lambda case: case.description, +) +def test_given_native_and_expansion_failures_when_linting_then_first_in_file_order_is_raised( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: BatchedPreparationFailureTestCase +) -> None: + write_lint_project( + root=tmp_path, + files={ + "models/customers.sql": _CUSTOMERS_MODEL, + "models/orders.sql": _ORDERS_MODEL, + test_case.unexpandable_file: 'MODEL (description "Broken");\nSELECT @unknown()\n', + }, + ) + monkeypatch.setattr(native_module, "prepare_lint_sql_batch", prepare_nothing_in_batch) + monkeypatch.setattr(native_module, "prepare_lint_sql", refuse_preparation) + + with pytest.raises(Exception) as raised: + _ = run_lint( + project_dir=tmp_path, + config=LintConfig(dialect="duckdb"), + skip_unexpandable=test_case.skip_unexpandable, + ) + + assert type(raised.value) is test_case.expected_error + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rule_reads.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rule_reads.py index a0fe767bcf..617e9cf8b9 100644 --- a/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rule_reads.py +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/engine/test_custom_rule_reads.py @@ -245,14 +245,14 @@ def test_given_file_edited_while_rules_run_when_rerunning_then_new_contents_are_ module_import=test_case.module_prelude, ) project: CompiledProject = two_model_project(customers_sql=_CUSTOMERS_SQL, customers_config={}) - run_host: Callable[..., object] = custom_rules._run_host + run_host: Callable[..., object] = custom_rules.run_custom_hosts def run_host_then_edit(**kwargs: Any) -> object: result: object = run_host(**kwargs) owners.write_text(_OWNERS_AFTER, encoding="utf-8") return result - monkeypatch.setattr(custom_rules, "_run_host", run_host_then_edit) + monkeypatch.setattr(custom_rules, "run_custom_hosts", run_host_then_edit) during_edit: CustomRulesOutcome = evaluate_cached_custom_rules( project=project, project_dir=tmp_path, cache_enabled=True ) diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py new file mode 100644 index 0000000000..6ae0c71543 --- /dev/null +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/_test_types.py @@ -0,0 +1,60 @@ +"""Test case declarations for custom-rule host partitioning.""" + +from dataclasses import dataclass + + +@dataclass(frozen=True) +class HostPartitionTestCase: + """One partitioned custom-rule evaluation compared with a single host.""" + + description: str + model_count: int + hosts: int + expected_host_runs: int + records_module_state: bool = False + detects_duplicates: bool = False + cache_enabled: bool = True + + +@dataclass(frozen=True) +class HostFailureTestCase: + """Failing subjects whose first single-host failure must be reported.""" + + description: str + model_count: int + failing_models: tuple[str, ...] + expected_failing_model: str + + +@dataclass(frozen=True) +class HostPlanTestCase: + """One deterministic split of planned invocations.""" + + description: str + plan: dict[str, list[str] | None] + invocations: tuple[tuple[str, str], ...] + hosts: int + expected_plans: tuple[dict[str, list[str] | None], ...] + + +@dataclass(frozen=True) +class HostCancellationTestCase: + """A fast failure on one host while the other hosts are still busy.""" + + description: str + model_count: int + hosts: int + timeout_millis: int + delay_seconds: float + expected_error_fragment: str + expected_max_seconds: float + + +@dataclass(frozen=True) +class HostCapacityTestCase: + """Fake cgroup files and the CPU count they allow.""" + + description: str + proc_cgroup: str + files: dict[str, str] + expected_cores: int diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py new file mode 100644 index 0000000000..cf8d7e815b --- /dev/null +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/helpers.py @@ -0,0 +1,138 @@ +"""Builders for custom-rule host partitioning tests.""" + +import os +from collections.abc import Callable +from dataclasses import replace +from operator import attrgetter +from pathlib import Path + +import pytest + +import sqlbuild._native as native_module +from sqlbuild.compiler.compile.models import CompiledProject +from sqlbuild.rule_engine._helpers.engine.catalogue import build_catalogue +from sqlbuild.rule_engine._helpers.engine.custom_rules import evaluate_custom_rules_cached +from sqlbuild.rule_engine._helpers.host import custom_host_pool +from sqlbuild.rule_engine.models import CustomRulesOutcome, Rule, RulesCacheConfig, RulesConfig +from tests.unit.src.sqlbuild.rule_engine.main.evaluate.helpers import build_project + +_MODEL_RULE: str = "XSQBRT101" +_PROJECT_RULE: str = "XSQBRT102" +_RULES_MODULE: str = f'''from sqlbuild.rules import Finding, Model, Project, RuleContext, rule + + +@rule(code="{_MODEL_RULE}", message="order models need an even suffix", remediation="Rename it.") +def even_suffix(*, model: Model, ctx: RuleContext) -> list[Finding]: + if model.name in FAILING: + raise ValueError(f"cannot inspect {{model.name}}") + time.sleep(DELAY_SECONDS) + if RECORD_SEEN: + SEEN[model.name] = True + source: str = ctx.sql.for_model(model).expanded.source + parity: str = "odd" if "odd" in source else "even" + if DETECT_DUPLICATES: + duplicate: bool = parity in SEEN + SEEN[parity] = True + return [ctx.finding(subject=model)] if duplicate else [] + return [ctx.finding(subject=model)] if parity == "odd" else [] + + +@rule(code="{_PROJECT_RULE}", message="projects need few models", remediation="Split it.") +def few_models(*, project: Project, ctx: RuleContext) -> list[Finding]: + del project + first: Model = ctx.project.models[0] + return [ctx.finding(subject=first)] if len(ctx.project.models) > 3 else [] +''' +_PARITIES: tuple[str, str] = ("even", "odd") + + +def write_order_rules( + *, + root: Path, + failing: tuple[str, ...] = (), + records_module_state: bool = False, + detects_duplicates: bool = False, + delay_seconds: float = 0.0, +) -> None: + """Write one model rule and one project rule with optional failures, state, or delays.""" + + rules: Path = root / "rules" / "orders.py" + rules.parent.mkdir(parents=True, exist_ok=True) + rules.write_text( + f"import time\n\nFAILING = {set(failing)!r}\nRECORD_SEEN = {records_module_state!r}\n" + f"DETECT_DUPLICATES = {detects_duplicates!r}\nDELAY_SECONDS = {delay_seconds!r}\n" + f"SEEN: dict[str, bool] = {{}}\n{_RULES_MODULE}", + encoding="utf-8", + ) + + +def orders_project(*, model_count: int, edits: dict[str, str] | None = None) -> CompiledProject: + """Return numbered order models whose SQL marks them odd or even, with optional edits.""" + + sql_by_name: dict[str, str] = { + f"orders_{index:03d}": f"SELECT 1 AS order_id -- {_PARITIES[index % 2]}" + for index in range(model_count) + } + sql_by_name.update(edits or {}) + projects: list[CompiledProject] = [ + build_project(name=name, relative_path=f"models/{name}.sql", sql=sql, config_values={}) + for name, sql in sql_by_name.items() + ] + return replace(projects[0], models=tuple(project.models[0] for project in projects)) + + +def evaluate_order_rules( + *, project: CompiledProject, root: Path, cache_enabled: bool = True +) -> CustomRulesOutcome: + """Evaluate both order rules through the incremental custom-rule cache.""" + + config: RulesConfig = RulesConfig( + select=(_MODEL_RULE, _PROJECT_RULE), cache=RulesCacheConfig(enabled=cache_enabled) + ) + rules: tuple[Rule, ...] = tuple( + filter(attrgetter("custom"), build_catalogue(config=config, project_dir=root)) + ) + return evaluate_custom_rules_cached( + project=project, config=config, project_dir=root, rules=rules, dialect="duckdb" + ) + + +def use_hosts(*, monkeypatch: pytest.MonkeyPatch, hosts: int) -> list[str]: + """Allow up to `hosts` host processes for any plan and record every host launch.""" + + launches: list[str] = [] + launch: Callable[[str], str] = native_module.run_custom_host_json + + def recording_launch(spec_json: str) -> str: + launches.append(spec_json) + return launch(spec_json) + + monkeypatch.setattr(custom_host_pool, "_MIN_INVOCATIONS_PER_HOST", 1) + monkeypatch.setattr(custom_host_pool, "available_cores", lambda: hosts) + monkeypatch.setattr(native_module, "run_custom_host_json", recording_launch) + return launches + + +def write_cgroup_files(*, root: Path, proc_cgroup: str, files: dict[str, str]) -> Path: + """Write a fake `/proc/self/cgroup` and cgroup hierarchy beneath `root`.""" + + for relative_path, contents in files.items(): + path: Path = root / "cgroup" / relative_path + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(contents, encoding="utf-8", errors="surrogateescape") + proc: Path = root / "proc-self-cgroup" + proc.write_text(proc_cgroup, encoding="utf-8", errors="surrogateescape") + return proc + + +def record_signals(*, monkeypatch: pytest.MonkeyPatch) -> list[tuple[int, int]]: + """Replace process signalling with a recorder so no process can be signalled.""" + + signals: list[tuple[int, int]] = [] + + def record(pid: int, signal_number: int) -> None: + signals.append((pid, signal_number)) + + monkeypatch.setattr(os, "kill", record) + monkeypatch.setattr(os, "killpg", record, raising=False) + return signals 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 new file mode 100644 index 0000000000..7f1eccf405 --- /dev/null +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_custom_host_pool.py @@ -0,0 +1,284 @@ +"""Custom rules split across host processes report exactly what one host reports.""" + +import time +from pathlib import Path + +import pytest + +from sqlbuild.compiler.compile.models import CompiledProject +from sqlbuild.rule_engine._helpers.host import custom_host_pool +from sqlbuild.rule_engine.constants import CUSTOM_RULES_CACHE_FILE +from sqlbuild.rule_engine.exceptions import RulesError +from sqlbuild.rule_engine.models import CustomRulesOutcome +from tests.unit.src.sqlbuild.rule_engine._helpers.host._test_types import ( + HostCancellationTestCase, + HostFailureTestCase, + HostPartitionTestCase, + HostPlanTestCase, +) +from tests.unit.src.sqlbuild.rule_engine._helpers.host.helpers import ( + evaluate_order_rules, + orders_project, + record_signals, + use_hosts, + write_order_rules, +) + + +@pytest.mark.parametrize( + "test_case", + [ + HostPartitionTestCase( + description="three hosts", model_count=7, hosts=3, expected_host_runs=3 + ), + HostPartitionTestCase( + description="more hosts than subjects", model_count=2, hosts=8, expected_host_runs=3 + ), + HostPartitionTestCase( + description="module state reruns the rule in one host", + model_count=7, + hosts=3, + expected_host_runs=4, + records_module_state=True, + ), + HostPartitionTestCase( + description="cross-model duplicate detector reruns in one host", + model_count=7, + hosts=3, + expected_host_runs=4, + detects_duplicates=True, + ), + HostPartitionTestCase( + description="uncached cross-model duplicate detector reruns in one host", + model_count=7, + hosts=3, + expected_host_runs=4, + detects_duplicates=True, + cache_enabled=False, + ), + ], + ids=lambda case: case.description, +) +def test_given_partitioned_hosts_when_evaluating_then_findings_and_cache_match_one_host( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostPartitionTestCase +) -> None: + write_order_rules( + root=tmp_path, + records_module_state=test_case.records_module_state, + detects_duplicates=test_case.detects_duplicates, + ) + project: CompiledProject = orders_project(model_count=test_case.model_count) + cache_path: Path = tmp_path / CUSTOM_RULES_CACHE_FILE + cache_glob: str = CUSTOM_RULES_CACHE_FILE + cache_enabled: bool = test_case.cache_enabled + + single_launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=1) + single: CustomRulesOutcome = evaluate_order_rules( + project=project, root=tmp_path, cache_enabled=cache_enabled + ) + single_warm: CustomRulesOutcome = evaluate_order_rules( + project=project, root=tmp_path, cache_enabled=cache_enabled + ) + single_cache: bytes = b"".join(path.read_bytes() for path in tmp_path.glob(cache_glob)) + cache_path.unlink(missing_ok=True) + monkeypatch.undo() + split_launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) + split: CustomRulesOutcome = evaluate_order_rules( + project=project, root=tmp_path, cache_enabled=cache_enabled + ) + split_cold_launches: int = len(split_launches) + split_cache: bytes = b"".join(path.read_bytes() for path in tmp_path.glob(cache_glob)) + split_warm: CustomRulesOutcome = evaluate_order_rules( + project=project, root=tmp_path, cache_enabled=cache_enabled + ) + + assert split.findings == single.findings + assert split_warm.findings == single_warm.findings == single.findings + assert len(single_launches) == 1 + int(not cache_enabled) + assert split_cold_launches == test_case.expected_host_runs + assert (split.cache_hits, split.cache_misses) == (single.cache_hits, single.cache_misses) + assert (split_warm.cache_hits, split_warm.cache_misses) == ( + single_warm.cache_hits, + single_warm.cache_misses, + ) + assert split_cache == single_cache + assert bool(single_cache) is cache_enabled + assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) + + +@pytest.mark.parametrize( + "test_case", + [ + HostPartitionTestCase( + description="edit after split", model_count=6, hosts=3, expected_host_runs=3 + ) + ], + ids=lambda case: case.description, +) +def test_given_results_cached_by_split_hosts_when_rerunning_then_only_edited_subjects_rerun( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostPartitionTestCase +) -> None: + write_order_rules(root=tmp_path) + project: CompiledProject = orders_project(model_count=test_case.model_count) + edited: CompiledProject = orders_project( + model_count=test_case.model_count, edits={"orders_003": "SELECT 2 AS order_id"} + ) + launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) + + cold: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) + cold_launches: int = len(launches) + warm: CustomRulesOutcome = evaluate_order_rules(project=project, root=tmp_path) + warm_launches: int = len(launches) - cold_launches + after_edit: CustomRulesOutcome = evaluate_order_rules(project=edited, root=tmp_path) + oracle: CustomRulesOutcome = evaluate_order_rules( + project=edited, root=tmp_path, cache_enabled=False + ) + + assert cold_launches == test_case.expected_host_runs + assert warm_launches == 0 + assert (warm.cache_hits, warm.cache_misses) == (cold.cache_misses, 0) + assert warm.findings == cold.findings + assert (after_edit.cache_hits, after_edit.cache_misses) == (cold.cache_misses - 1, 1) + assert after_edit.findings == oracle.findings + + +@pytest.mark.parametrize( + "test_case", + [ + HostFailureTestCase( + description="earliest failure lives on a later host", + model_count=6, + failing_models=("orders_001", "orders_003"), + expected_failing_model="orders_001", + ), + HostFailureTestCase( + description="one failing subject", + model_count=6, + failing_models=("orders_005",), + expected_failing_model="orders_005", + ), + ], + ids=lambda case: case.description, +) +def test_given_failing_rule_on_split_hosts_when_evaluating_then_single_host_error_is_raised( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostFailureTestCase +) -> None: + write_order_rules(root=tmp_path, failing=test_case.failing_models) + project: CompiledProject = orders_project(model_count=test_case.model_count) + + _ = use_hosts(monkeypatch=monkeypatch, hosts=1) + with pytest.raises(RulesError) as single: + _ = evaluate_order_rules(project=project, root=tmp_path) + monkeypatch.undo() + _ = use_hosts(monkeypatch=monkeypatch, hosts=3) + with pytest.raises(RulesError) as split: + _ = evaluate_order_rules(project=project, root=tmp_path) + + assert str(split.value) == str(single.value) + assert f"cannot inspect {test_case.expected_failing_model}" in str(split.value) + + +@pytest.mark.parametrize( + "test_case", + [ + HostPlanTestCase( + description="project rules and whole subjects are dealt round-robin", + plan={"XSQBRT101": None, "XSQBRT102": None, "XSQBRT103": ["models/b.sql"]}, + invocations=( + ("XSQBRT102", ""), + ("XSQBRT101", "models/a.sql"), + ("XSQBRT101", "models/b.sql"), + ("XSQBRT103", "models/b.sql"), + ("XSQBRT101", "models/c.sql"), + ), + hosts=2, + expected_plans=( + {"XSQBRT101": ["models/b.sql"], "XSQBRT102": [""], "XSQBRT103": ["models/b.sql"]}, + {"XSQBRT101": ["models/a.sql", "models/c.sql"], "XSQBRT102": [], "XSQBRT103": []}, + ), + ), + HostPlanTestCase( + description="one subject keeps the original plan", + plan={"XSQBRT101": None}, + invocations=(("XSQBRT101", "models/a.sql"),), + hosts=4, + expected_plans=({"XSQBRT101": None},), + ), + ], + ids=lambda case: case.description, +) +def test_given_planned_invocations_when_partitioning_then_split_is_deterministic( + test_case: HostPlanTestCase, +) -> None: + plans: tuple[dict[str, list[str] | None], ...] = custom_host_pool.partition_host_plan( + plan=test_case.plan, invocations=test_case.invocations, hosts=test_case.hosts + ) + + assert plans == test_case.expected_plans + + +@pytest.mark.parametrize( + "test_case", + [ + HostCancellationTestCase( + description="busy hosts stop at their next invocation", + model_count=60, + hosts=3, + timeout_millis=120_000, + delay_seconds=2.0, + expected_error_fragment="cannot inspect orders_000", + expected_max_seconds=20.0, + ), + ], + ids=lambda case: case.description, +) +def test_given_failing_host_when_others_are_busy_then_they_stop_before_later_invocations( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostCancellationTestCase +) -> None: + write_order_rules(root=tmp_path, failing=("orders_000",), delay_seconds=test_case.delay_seconds) + project: CompiledProject = orders_project(model_count=test_case.model_count) + launches: list[str] = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) + monkeypatch.setattr(custom_host_pool, "_HOST_TIMEOUT_MILLIS", test_case.timeout_millis) + started: float = time.monotonic() + + with pytest.raises(RulesError) as raised: + _ = evaluate_order_rules(project=project, root=tmp_path) + + assert time.monotonic() - started < test_case.expected_max_seconds + assert test_case.expected_error_fragment in str(raised.value) + assert len(launches) == test_case.hosts + 1 + assert not tuple((tmp_path / "target" / "rules-cache" / "host-inputs").iterdir()) + + +@pytest.mark.parametrize( + "test_case", + [ + HostCancellationTestCase( + description="busy hosts stop at their next invocation", + model_count=60, + hosts=3, + timeout_millis=120_000, + delay_seconds=2.0, + expected_error_fragment="cannot inspect orders_000", + expected_max_seconds=20.0, + ), + ], + ids=lambda case: case.description, +) +def test_given_cancelled_hosts_when_failing_then_no_process_is_signalled( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostCancellationTestCase +) -> None: + write_order_rules(root=tmp_path, failing=("orders_000",), delay_seconds=test_case.delay_seconds) + project: CompiledProject = orders_project(model_count=test_case.model_count) + _ = use_hosts(monkeypatch=monkeypatch, hosts=test_case.hosts) + signals: list[tuple[int, int]] = record_signals(monkeypatch=monkeypatch) + + with pytest.raises(RulesError) as raised: + _ = evaluate_order_rules(project=project, root=tmp_path) + + assert signals == [] + assert test_case.expected_error_fragment in str(raised.value) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_host_capacity.py b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_host_capacity.py new file mode 100644 index 0000000000..025e8f4757 --- /dev/null +++ b/tests/unit/src/sqlbuild/rule_engine/_helpers/host/test_host_capacity.py @@ -0,0 +1,134 @@ +"""Host capacity honours container CPU quotas from cgroup v1 and v2.""" + +import os +from pathlib import Path + +import pytest + +from sqlbuild.rule_engine._helpers.host import custom_host_pool +from sqlbuild.rule_engine._helpers.host.host_capacity import available_cores +from tests.unit.src.sqlbuild.rule_engine._helpers.host._test_types import HostCapacityTestCase +from tests.unit.src.sqlbuild.rule_engine._helpers.host.helpers import write_cgroup_files + +_NODE_CPUS: int = 64 + + +@pytest.mark.parametrize( + "test_case", + [ + HostCapacityTestCase( + description="cgroup v2 quota on the process cgroup", + proc_cgroup="0::/kubepods/orders\n", + files={"kubepods/orders/cpu.max": "400000 100000\n"}, + expected_cores=4, + ), + HostCapacityTestCase( + description="cgroup v2 fractional quota rounds up", + proc_cgroup="0::/kubepods/orders\n", + files={"kubepods/orders/cpu.max": "150000 100000\n"}, + expected_cores=2, + ), + HostCapacityTestCase( + description="cgroup v2 ancestor quota is tighter", + proc_cgroup="0::/kubepods/orders\n", + files={ + "kubepods/cpu.max": "200000 100000\n", + "kubepods/orders/cpu.max": "800000 100000\n", + }, + expected_cores=2, + ), + HostCapacityTestCase( + description="cgroup v2 without a quota", + proc_cgroup="0::/\n", + files={"cpu.max": "max 100000\n"}, + expected_cores=_NODE_CPUS, + ), + HostCapacityTestCase( + description="cgroup v1 CFS quota", + proc_cgroup="4:cpu,cpuacct:/docker/orders\n", + files={ + "cpu,cpuacct/docker/orders/cpu.cfs_quota_us": "300000\n", + "cpu,cpuacct/docker/orders/cpu.cfs_period_us": "100000\n", + }, + expected_cores=3, + ), + HostCapacityTestCase( + description="cgroup v1 unlimited CFS quota", + proc_cgroup="4:cpu,cpuacct:/\n", + files={ + "cpu,cpuacct/cpu.cfs_quota_us": "-1\n", + "cpu,cpuacct/cpu.cfs_period_us": "100000\n", + }, + expected_cores=_NODE_CPUS, + ), + HostCapacityTestCase( + description="no cgroup files", proc_cgroup="", files={}, expected_cores=_NODE_CPUS + ), + HostCapacityTestCase( + description="non-UTF-8 cgroup v2 quota is ignored", + proc_cgroup="0::/kubepods/orders\n", + files={"kubepods/orders/cpu.max": "\udcff\udcfe 100000\n"}, + expected_cores=_NODE_CPUS, + ), + HostCapacityTestCase( + description="non-UTF-8 process cgroup is ignored", + proc_cgroup="0::/kubepods/\udcff\n", + files={"kubepods/cpu.max": "200000 100000\n"}, + expected_cores=_NODE_CPUS, + ), + HostCapacityTestCase( + description="non-UTF-8 cgroup v1 quota is ignored", + proc_cgroup="4:cpu,cpuacct:/\n", + files={ + "cpu,cpuacct/cpu.cfs_quota_us": "\udcff\n", + "cpu,cpuacct/cpu.cfs_period_us": "100000\n", + }, + expected_cores=_NODE_CPUS, + ), + ], + ids=lambda case: case.description, +) +def test_given_cgroup_quota_when_counting_cores_then_quota_caps_process_cpus( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, test_case: HostCapacityTestCase +) -> None: + proc_cgroup: Path = write_cgroup_files( + root=tmp_path, proc_cgroup=test_case.proc_cgroup, files=test_case.files + ) + monkeypatch.setattr(os, "process_cpu_count", lambda: _NODE_CPUS, raising=False) + + cores: int = available_cores(cgroup_root=tmp_path / "cgroup", proc_cgroup=proc_cgroup) + + assert cores == test_case.expected_cores + + +@pytest.mark.parametrize( + "test_case", + [ + HostCapacityTestCase( + description="hosts are capped on a large node", + proc_cgroup="", + files={}, + expected_cores=8, + ) + ], + ids=lambda case: case.description, +) +def test_given_many_cores_when_partitioning_then_at_most_eight_hosts_start( + monkeypatch: pytest.MonkeyPatch, test_case: HostCapacityTestCase +) -> None: + monkeypatch.setattr(custom_host_pool, "available_cores", lambda: _NODE_CPUS) + invocations: tuple[tuple[str, str], ...] = tuple( + ("XSQBRT101", f"models/orders_{index:04d}.sql") for index in range(10_000) + ) + + plans: tuple[dict[str, list[str] | None], ...] = custom_host_pool.partition_host_plan( + plan={"XSQBRT101": None}, + invocations=invocations, + hosts=custom_host_pool.host_count(len(invocations)), + ) + + assert len(plans) == test_case.expected_cores + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/website/src/content/docs/docs/concepts/rules/custom-rules/testing-and-determinism.mdx b/website/src/content/docs/docs/concepts/rules/custom-rules/testing-and-determinism.mdx index aa6facc605..e0c011bbdc 100644 --- a/website/src/content/docs/docs/concepts/rules/custom-rules/testing-and-determinism.mdx +++ b/website/src/content/docs/docs/concepts/rules/custom-rules/testing-and-determinism.mdx @@ -37,6 +37,19 @@ the next. Compute findings from the subject and `ctx` alone. Repository pytest files remain outside the SQLBuild project's SQL `tests/` directory. +## Parallel evaluation + +On large projects, SQLBuild splits custom Rule evaluation across several host processes, at most +one per available CPU and never more than eight. Each process evaluates a share of the subjects, +and SQLBuild merges the findings back into the order a single process reports them. + +A Rule that changes module-level state while it runs, for example by adding each model to a set of +already-seen names, could report different findings when its subjects are split. SQLBuild detects +the state change and evaluates that Rule again in one process over all of its subjects, so its +findings and cache entries match a single-process run. The result stays correct, but the Rule no +longer runs in parallel and costs an extra host process. Keep each subject's findings independent +of the others so Rules can run in parallel. + ## Helpers Custom Rules can import repository-owned helpers under `rules/` and these deterministic standard