From c289dad24d464555345d0f5fbd48573d58a7e4fe Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sat, 3 Oct 2026 03:57:30 +0100 Subject: [PATCH 1/4] perf(compile): publish staged cold artifacts by rename instead of copying Background artifact preparation wrote every compiled file to a system temp directory and the write phase then walked, read and rewrote each one into target/. Stage beside target/ on the same filesystem instead, move the whole tree into place when target/compiled is absent, and otherwise rename new files and rewrite only changed ones, keeping the unchanged-file and stale-file semantics. Each staging directory is claimed by locking a sibling owner lock before the directory exists, held for the whole compile. At the start of a compile, staging directories are removed only while this compile holds their owner lock for the entire removal; cleanup stops and logs instead of continuing after a failed removal. Publication verifies the owner marker and the staged file set before moving anything or deleting stale files, and fails the compile rather than deleting valid artifacts if staging changed. New files that cannot be renamed across filesystems (EXDEV, for example target/compiled symlinked to another mount) are copied through a sibling temporary file and replaced atomically. --- .../cli/commands/_helpers/compile/pipeline.py | 4 +- .../_helpers/compile/target_writer.py | 126 ++++++++-- .../classes/prepared_compile_artifacts.py | 129 ++++++++-- src/sqlbuild/cli/commands/exceptions.py | 4 + .../cli/commands/main/project/_compile.py | 1 + .../sqlbuild/cli/commands/main/_test_types.py | 20 +- .../src/sqlbuild/cli/commands/main/helpers.py | 50 +++- .../main/test_compile_prepared_artifacts.py | 102 +++++++- .../cli/commands/main/compile/_test_types.py | 23 ++ .../cli/commands/main/compile/helpers.py | 35 +++ .../test_staged_artifact_publication.py | 226 ++++++++++++++++++ 11 files changed, 675 insertions(+), 45 deletions(-) create mode 100644 tests/unit/src/sqlbuild/cli/commands/main/compile/test_staged_artifact_publication.py diff --git a/src/sqlbuild/cli/commands/_helpers/compile/pipeline.py b/src/sqlbuild/cli/commands/_helpers/compile/pipeline.py index 76982efdca..72b2548e2a 100644 --- a/src/sqlbuild/cli/commands/_helpers/compile/pipeline.py +++ b/src/sqlbuild/cli/commands/_helpers/compile/pipeline.py @@ -206,7 +206,9 @@ def _analyze_compile_project( and rules_config.select and not any(item.is_error for item in core_diagnostics) ): - prepared_artifacts.start(project=graph.project, adapter=adapter) + prepared_artifacts.start( + project=graph.project, adapter=adapter, target_dir=project_dir / "target" + ) _ = start_compile_phase(status=status, message="Evaluating built-in and custom rules...") rules_start: float = time.monotonic() rules_result = run_rules( diff --git a/src/sqlbuild/cli/commands/_helpers/compile/target_writer.py b/src/sqlbuild/cli/commands/_helpers/compile/target_writer.py index 583da81cc6..0b07b618f5 100644 --- a/src/sqlbuild/cli/commands/_helpers/compile/target_writer.py +++ b/src/sqlbuild/cli/commands/_helpers/compile/target_writer.py @@ -2,8 +2,10 @@ from __future__ import annotations +import errno import json import os +import uuid from pathlib import Path from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter @@ -16,6 +18,7 @@ sql_test_artifact_record_key, write_sql_test_artifact_cache, ) +from sqlbuild.cli.commands.exceptions import StagedArtifactsChangedError from sqlbuild.cli.compile.models import ( SqlTestArtifactCacheRecord, SqlTestArtifactIdentityContext, @@ -168,27 +171,45 @@ def write_static_compile_target( ) +def staged_artifact_files(*, target_dir: Path) -> frozenset[str]: + """Return the relative paths of every compiled artifact staged under target_dir.""" + + return frozenset( + relative for _, relative in _staged_files(staged_dir=target_dir / _COMPILED_DIR) + ) + + def publish_static_compile_target( - *, prepared: WrittenTarget, target_dir: Path, manifest: dict[str, object] | None + *, + prepared: WrittenTarget, + target_dir: Path, + manifest: dict[str, object] | None, + expected_files: frozenset[str], ) -> WrittenTarget: """Publish staged files with the same unchanged-file and stale-file semantics.""" compiled_dir: Path = target_dir / _COMPILED_DIR staged_dir: Path = prepared.target_dir / _COMPILED_DIR - check_existing: bool = compiled_dir.is_dir() - managed_paths: set[Path] = set() - target_dir.mkdir(parents=True, exist_ok=True) - for source in sorted(staged_dir.rglob("*")): - if not source.is_file(): - continue - path: Path = compiled_dir / source.relative_to(staged_dir) - _write_bytes_if_changed( - path=path, contents=source.read_bytes(), check_existing=check_existing + staged_files: list[tuple[Path, str]] = _staged_files(staged_dir=staged_dir) + if frozenset(relative for _, relative in staged_files) != expected_files: + raise StagedArtifactsChangedError( + "staged compile artifacts changed before publication; no artifact was published", + help="Rerun the compile; avoid deleting target/ while a compile is running.", ) - managed_paths.add(path) - if check_existing: - with record_compile_timing("stale_traversal_ms"): - _remove_stale_compiled_files(target_dir=target_dir, managed_paths=managed_paths) + target_dir.mkdir(parents=True, exist_ok=True) + moved: bool = not compiled_dir.is_dir() and _move_staged_tree( + staged_dir=staged_dir, path=compiled_dir + ) + if not moved: + check_existing: bool = compiled_dir.is_dir() + managed_paths: set[Path] = set() + for source, relative in staged_files: + path: Path = compiled_dir / relative + _publish_staged_file(source=source, path=path, check_existing=check_existing) + managed_paths.add(path) + if check_existing: + with record_compile_timing("stale_traversal_ms"): + _remove_stale_compiled_files(target_dir=target_dir, managed_paths=managed_paths) if manifest is not None: _write_manifest(target_dir=target_dir, manifest=manifest) return WrittenTarget( @@ -202,6 +223,67 @@ def publish_static_compile_target( ) +def _move_staged_tree(*, staged_dir: Path, path: Path) -> bool: + with record_compile_timing("physical_write_ms"): + try: + os.rename(staged_dir, path) + except OSError: + return False + return True + + +def _staged_files(*, staged_dir: Path) -> list[tuple[Path, str]]: + root_prefix_length: int = len(os.fspath(staged_dir)) + 1 + files: list[tuple[Path, str]] = [] + for root, _, filenames in os.walk(staged_dir): + relative_root: str = root[root_prefix_length:] + for filename in filenames: + files.append((Path(root, filename), os.path.join(relative_root, filename))) + return files + + +def _publish_staged_file(*, source: Path, path: Path, check_existing: bool) -> None: + with record_compile_timing("physical_write_ms"): + if check_existing and path.is_file(): + contents: bytes = source.read_bytes() + existing: bytes = path.read_bytes() + if existing == contents: + return + _ = existing.decode("utf-8") + _overwrite_bytes(path=path, contents=contents) + return + try: + _move_new_file(source=source, path=path) + except FileNotFoundError: + path.parent.mkdir(parents=True, exist_ok=True) + _move_new_file(source=source, path=path) + + +def _move_new_file(*, source: Path, path: Path) -> None: + try: + os.replace(source, path) + except OSError as error: + if error.errno != errno.EXDEV: + raise + _copy_new_file(source=source, path=path) + + +def _copy_new_file(*, source: Path, path: Path) -> None: + """Copy across filesystems through a sibling temporary file so the new file is atomic.""" + + temporary: Path = path.with_name(f".{path.name}.{uuid.uuid4().hex}.tmp") + descriptor: int = os.open(temporary, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o666) + try: + try: + _write_all(descriptor=descriptor, path=path, contents=source.read_bytes()) + finally: + os.close(descriptor) + os.replace(temporary, path) + except BaseException: + temporary.unlink(missing_ok=True) + raise + + def _write_models(*, target_dir: Path, plan_output: PlanOutput, check_existing: bool) -> set[Path]: """Write model resolved SQL.""" @@ -568,16 +650,20 @@ def _overwrite_text(*, path: Path, contents: str) -> None: def _overwrite_bytes(*, path: Path, contents: bytes) -> None: descriptor: int = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o666) try: - offset: int = 0 - while offset < len(contents): - written: int = os.write(descriptor, contents[offset:]) - if written == 0: - raise OSError(f"failed to write compiled artifact '{path}'") - offset += written + _write_all(descriptor=descriptor, path=path, contents=contents) finally: os.close(descriptor) +def _write_all(*, descriptor: int, path: Path, contents: bytes) -> None: + offset: int = 0 + while offset < len(contents): + written: int = os.write(descriptor, contents[offset:]) + if written == 0: + raise OSError(f"failed to write compiled artifact '{path}'") + offset += written + + def _remove_stale_compiled_files(*, target_dir: Path, managed_paths: set[Path]) -> None: compiled_dir: Path = target_dir / _COMPILED_DIR if not compiled_dir.is_dir(): diff --git a/src/sqlbuild/cli/commands/classes/prepared_compile_artifacts.py b/src/sqlbuild/cli/commands/classes/prepared_compile_artifacts.py index 0dd6f0f036..473b1213ed 100644 --- a/src/sqlbuild/cli/commands/classes/prepared_compile_artifacts.py +++ b/src/sqlbuild/cli/commands/classes/prepared_compile_artifacts.py @@ -2,22 +2,33 @@ from __future__ import annotations +import logging +import shutil +import uuid from concurrent.futures import Future, ThreadPoolExecutor from contextvars import Context, copy_context from pathlib import Path -from tempfile import TemporaryDirectory from types import TracebackType +from filelock import FileLock, Timeout + from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter from sqlbuild.cli.commands._helpers.compile.target_writer import ( publish_static_compile_target, + staged_artifact_files, write_static_compile_target, ) +from sqlbuild.cli.commands.exceptions import StagedArtifactsChangedError from sqlbuild.cli.output.models import WrittenTarget from sqlbuild.compiler.compile.models import CompiledProject, CompilerDiagnostic _MIN_PREPARED_ARTIFACT_MODELS: int = 128 _MAX_PREPARED_ARTIFACT_MODELS: int = 5000 +_STAGING_PREFIX: str = ".sqlbuild-staging-" +_LOCK_SUFFIX: str = ".lock" +_OWNER_MARKER: str = ".owner" +_LOGGER: logging.Logger = logging.getLogger(__name__) +type _PreparedResult = tuple[WrittenTarget, frozenset[str]] class PreparedCompileArtifacts: @@ -25,9 +36,11 @@ class PreparedCompileArtifacts: def __init__(self, *, enabled: bool) -> None: self.enabled: bool = enabled - self._directory: TemporaryDirectory[str] | None = None + self._staging_dir: Path | None = None + self._owner_lock: FileLock | None = None + self._created_target_dir: Path | None = None self._executor: ThreadPoolExecutor | None = None - self._future: Future[WrittenTarget] | None = None + self._future: Future[_PreparedResult] | None = None def __enter__(self) -> PreparedCompileArtifacts: return self @@ -42,10 +55,27 @@ def __exit__( if self._executor is not None: self._executor.shutdown(wait=True, cancel_futures=True) finally: - if self._directory is not None: - self._directory.cleanup() + if self._staging_dir is not None: + shutil.rmtree(self._staging_dir, ignore_errors=True) + if self._owner_lock is not None: + _release_and_unlink(self._owner_lock) + if self._created_target_dir is not None: + _remove_if_empty(self._created_target_dir) + + def remove_abandoned_staging(self, *, target_dir: Path) -> None: + """Delete staging directories whose owner lock this compile can hold for the removal.""" + + if not self.enabled: + return + for directory in sorted(target_dir.glob(f"{_STAGING_PREFIX}*")): + if directory.suffix == _LOCK_SUFFIX: + continue + if not _remove_abandoned_directory(directory): + return + + def start(self, *, project: CompiledProject, adapter: BaseAdapter, target_dir: Path) -> None: + """Stage artifacts beside target_dir so publication can move them instead of copying.""" - def start(self, *, project: CompiledProject, adapter: BaseAdapter) -> None: if ( not self.enabled or project.compile_cache_dir is not None @@ -53,35 +83,51 @@ def start(self, *, project: CompiledProject, adapter: BaseAdapter) -> None: or len(project.models) > _MAX_PREPARED_ARTIFACT_MODELS ): return + staging_dir: Path = target_dir / f"{_STAGING_PREFIX}{uuid.uuid4().hex}" try: - self._directory = TemporaryDirectory(prefix="sqlbuild-compile-") + if not target_dir.is_dir(): + target_dir.mkdir(parents=True) + self._created_target_dir = target_dir + self._owner_lock = FileLock(_owner_lock_path(staging_dir)) + self._owner_lock.acquire() + _create_staging_directory(staging_dir) except OSError: return + self._staging_dir = staging_dir self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-artifacts") context: Context = copy_context() - target_dir: Path = Path(self._directory.name) - def prepare() -> WrittenTarget: - return context.run( + def prepare() -> _PreparedResult: + written: WrittenTarget = context.run( write_static_compile_target, - target_dir=target_dir, + target_dir=staging_dir, adapter=adapter, project=project, ) + return written, staged_artifact_files(target_dir=staging_dir) self._future = self._executor.submit(prepare) def publish( self, *, target_dir: Path, manifest: dict[str, object] | None ) -> WrittenTarget | None: - if self._future is None: + if self._future is None or self._staging_dir is None: return None try: - prepared: WrittenTarget = self._future.result() + prepared, expected_files = self._future.result() except OSError: return None + if not _owns_staging_directory(self._staging_dir): + raise StagedArtifactsChangedError( + f"staging directory '{self._staging_dir.name}' was removed or replaced before " + "publication; no artifact was published", + help="Rerun the compile; avoid deleting target/ while a compile is running.", + ) return publish_static_compile_target( - prepared=prepared, target_dir=target_dir, manifest=manifest + prepared=prepared, + target_dir=target_dir, + manifest=manifest, + expected_files=expected_files, ) def planning_diagnostics(self) -> tuple[CompilerDiagnostic, ...] | None: @@ -90,6 +136,59 @@ def planning_diagnostics(self) -> tuple[CompilerDiagnostic, ...] | None: if self._future is None: return None try: - return self._future.result().diagnostics + return self._future.result()[0].diagnostics except OSError: return None + + +def _owner_lock_path(staging_dir: Path) -> Path: + return staging_dir.with_name(f"{staging_dir.name}{_LOCK_SUFFIX}") + + +def _create_staging_directory(staging_dir: Path) -> None: + staging_dir.mkdir() + (staging_dir / _OWNER_MARKER).write_text(staging_dir.name, encoding="utf-8") + + +def _owns_staging_directory(staging_dir: Path) -> bool: + try: + return (staging_dir / _OWNER_MARKER).read_text(encoding="utf-8") == staging_dir.name + except OSError: + return False + + +def _remove_abandoned_directory(directory: Path) -> bool: + """Remove one ownerless staging directory while holding its lock; False stops cleanup.""" + + lock: FileLock = FileLock(_owner_lock_path(directory), timeout=0) + try: + lock.acquire() + except Timeout: + return True + except OSError as error: + _LOGGER.warning("skipped abandoned staging cleanup in %s: %s", directory.parent, error) + return False + try: + if directory.is_dir(): + shutil.rmtree(directory) + except OSError as error: + lock.release() + _LOGGER.warning("stopped abandoned staging cleanup at %s: %s", directory, error) + return False + _release_and_unlink(lock) + return True + + +def _release_and_unlink(lock: FileLock) -> None: + try: + Path(lock.lock_file).unlink(missing_ok=True) + except OSError: + pass + lock.release() + + +def _remove_if_empty(directory: Path) -> None: + try: + directory.rmdir() + except OSError: + pass diff --git a/src/sqlbuild/cli/commands/exceptions.py b/src/sqlbuild/cli/commands/exceptions.py index 0294f1403e..fd9bc6a624 100644 --- a/src/sqlbuild/cli/commands/exceptions.py +++ b/src/sqlbuild/cli/commands/exceptions.py @@ -15,6 +15,10 @@ def __init__(self, message: str, *, code: str | None = None, help: str | None = self.help = help +class StagedArtifactsChangedError(CliUserError): + """Staged compile artifacts changed before publication, so nothing stale is deleted.""" + + class QueryExecutionError(CliUserError): """The warehouse rejected or failed an ad hoc query.""" diff --git a/src/sqlbuild/cli/commands/main/project/_compile.py b/src/sqlbuild/cli/commands/main/project/_compile.py index 08bfa4e7c0..666180c987 100644 --- a/src/sqlbuild/cli/commands/main/project/_compile.py +++ b/src/sqlbuild/cli/commands/main/project/_compile.py @@ -81,6 +81,7 @@ def _run_compile_with_status( manifest: bool = request.manifest no_color: bool = request.no_color lineage_mode: CompileLineageMode = request.lineage_mode + prepared_artifacts.remove_abandoned_staging(target_dir=project_dir / "target") analysis: CompileAnalysis = analyze_compile_project( project_dir=project_dir, no_sql_validation=request.no_sql_validation, 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 9287a1e455..51001c35ed 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/_test_types.py @@ -2,7 +2,7 @@ from collections.abc import Callable from dataclasses import dataclass, field -from tempfile import TemporaryDirectory +from pathlib import Path from sqlbuild.compiler.planner.models import CursorOverrides @@ -27,7 +27,23 @@ class VariedCompileFixtureTestCase: class PreparedArtifactsCompileTestCase: description: str model_count: int - temporary_directory_factory: Callable[..., TemporaryDirectory[str]] + staging_directory_factory: Callable[[Path], None] + expected_exit_code: int = 0 + + +@dataclass(frozen=True) +class AbandonedStagingCompileTestCase: + description: str + model_count: int + abandoned: tuple[str, ...] + locked: tuple[str, ...] + expected_remaining: tuple[str, ...] + + +@dataclass(frozen=True) +class CrossDeviceArtifactsCompileTestCase: + description: str + model_count: int expected_exit_code: int = 0 diff --git a/tests/integration/src/sqlbuild/cli/commands/main/helpers.py b/tests/integration/src/sqlbuild/cli/commands/main/helpers.py index 900532fb1c..c2d82be487 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/helpers.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/helpers.py @@ -1,9 +1,10 @@ from __future__ import annotations import json +import os +import tempfile from collections.abc import Callable from pathlib import Path -from tempfile import TemporaryDirectory from typing import Any, cast import duckdb @@ -25,8 +26,51 @@ from sqlbuild.spec.contracts.models import SourceLocation -def unavailable_artifact_directory(*, prefix: str) -> TemporaryDirectory[str]: - raise OSError("temporary storage unavailable") +def staging_directories(target_dir: Path) -> list[str]: + """Return hidden entries left in target_dir, such as abandoned artifact staging.""" + + return sorted(path.name for path in target_dir.glob(".*")) + + +def write_staging_directory(*, target_dir: Path, name: str) -> Path: + """Write one artifact staging directory with a staged model file.""" + + directory: Path = target_dir / name + (directory / "compiled" / "models").mkdir(parents=True) + (directory / "compiled" / "models" / "orders.sql").write_text("SELECT 1\n", encoding="utf-8") + return directory + + +def write_prepared_artifacts_project(*, project_dir: Path, model_count: int) -> None: + """Write a DuckDB project large enough to stage compile artifacts in the background.""" + + (project_dir / "sqlbuild_project.toml").write_text( + 'name = "orders"\nadapter = "duckdb"\n[rules]\nselect = ["SQBRSQL021"]\n', + encoding="utf-8", + ) + models: Path = project_dir / "models" + models.mkdir() + for index in range(model_count): + (models / f"orders_{index:03}.sql").write_text( + "MODEL (description 'Test model.', materialized table, contract enforced, " + "columns (order_id (type INTEGER)));\n" + f"SELECT CAST({index} AS INTEGER) AS order_id\n", + encoding="utf-8", + ) + + +def has_second_filesystem(root: Path) -> bool: + """Return whether root is a writable directory on another device than the temp directory.""" + + return ( + root.is_dir() + and os.access(root, os.W_OK) + and root.stat().st_dev != Path(tempfile.gettempdir()).stat().st_dev + ) + + +def unavailable_staging_directory(staging_dir: Path) -> None: + raise OSError(f"staging storage unavailable for {staging_dir.name}") def write_compile_startup_project(project_dir: Path) -> None: diff --git a/tests/integration/src/sqlbuild/cli/commands/main/test_compile_prepared_artifacts.py b/tests/integration/src/sqlbuild/cli/commands/main/test_compile_prepared_artifacts.py index 6b230cdad2..20374275d6 100644 --- a/tests/integration/src/sqlbuild/cli/commands/main/test_compile_prepared_artifacts.py +++ b/tests/integration/src/sqlbuild/cli/commands/main/test_compile_prepared_artifacts.py @@ -1,26 +1,41 @@ """Real cold compilation preserves publication gates and unchanged artifact files.""" import json +import os +from contextlib import ExitStack from pathlib import Path from tempfile import TemporaryDirectory import duckdb import pytest +from filelock import FileLock from sqlbuild.cli.commands.classes import prepared_compile_artifacts from sqlbuild.cli.commands.main.entrypoint.entry import main from tests.integration.src.sqlbuild.cli.commands.main._test_types import ( + AbandonedStagingCompileTestCase, + CrossDeviceArtifactsCompileTestCase, PreparedArtifactsCompileTestCase, ) -from tests.integration.src.sqlbuild.cli.commands.main.helpers import unavailable_artifact_directory +from tests.integration.src.sqlbuild.cli.commands.main.helpers import ( + has_second_filesystem, + staging_directories, + unavailable_staging_directory, + write_prepared_artifacts_project, + write_staging_directory, +) + +_SECOND_FILESYSTEM: Path = Path("/dev/shm") @pytest.mark.parametrize( "test_case", ( - PreparedArtifactsCompileTestCase("cold staged artifacts", 128, TemporaryDirectory), PreparedArtifactsCompileTestCase( - "unavailable temporary storage", 128, unavailable_artifact_directory + "cold staged artifacts", 128, prepared_compile_artifacts._create_staging_directory + ), + PreparedArtifactsCompileTestCase( + "unavailable staging storage", 128, unavailable_staging_directory ), ), ids=lambda case: case.description, @@ -32,7 +47,7 @@ def test_given_existing_artifacts_when_rules_fail_then_staged_changes_are_not_pu monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr( - prepared_compile_artifacts, "TemporaryDirectory", test_case.temporary_directory_factory + prepared_compile_artifacts, "_create_staging_directory", test_case.staging_directory_factory ) config_path: Path = tmp_path / "sqlbuild_project.toml" config: str = 'name = "orders"\nadapter = "duckdb"\n[rules]\nselect = ["SQBRSQL021"]\n' @@ -59,6 +74,7 @@ def test_given_existing_artifacts_when_rules_fail_then_staged_changes_are_not_pu before: dict[Path, bytes] = {path: path.read_bytes() for path in compiled.rglob("*.sql")} modified_times: dict[Path, int] = {path: path.stat().st_mtime_ns for path in before} assert len(before) == test_case.model_count + assert staging_directories(tmp_path / "target") == [] stale: Path = compiled / "models" / "obsolete.sql" stale.write_text("SELECT 0\n", encoding="utf-8") (models / "orders_000.sql").write_text( @@ -81,6 +97,7 @@ def test_given_existing_artifacts_when_rules_fail_then_staged_changes_are_not_pu assert failed["has_errors"] is True assert {path: path.read_bytes() for path in before} == before assert stale.exists() + assert staging_directories(tmp_path / "target") == [] config_path.write_text(config, encoding="utf-8") assert main(arguments) == 0 @@ -91,9 +108,86 @@ def test_given_existing_artifacts_when_rules_fail_then_staged_changes_are_not_pu for path in modified_times.keys() - {changed}: assert path.stat().st_mtime_ns == modified_times[path] assert path.read_bytes() == before[path] + assert staging_directories(tmp_path / "target") == [] with duckdb.connect() as connection: assert connection.execute(changed.read_text(encoding="utf-8")).fetchall() == [(2,)] +@pytest.mark.parametrize( + "test_case", + ( + AbandonedStagingCompileTestCase( + description="only staging whose owner lock is free is removed", + model_count=128, + abandoned=(".sqlbuild-staging-killed",), + locked=(".sqlbuild-staging-running",), + expected_remaining=(".sqlbuild-staging-running", ".sqlbuild-staging-running.lock"), + ), + ), + ids=lambda case: case.description, +) +def test_given_leftover_staging_when_compiling_then_only_abandoned_staging_is_removed( + test_case: AbandonedStagingCompileTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + write_prepared_artifacts_project(project_dir=tmp_path, model_count=test_case.model_count) + target: Path = tmp_path / "target" + for name in test_case.abandoned: + write_staging_directory(target_dir=target, name=name) + with ExitStack() as owners: + for name in test_case.locked: + owners.enter_context(FileLock(target / f"{name}.lock")) + write_staging_directory(target_dir=target, name=name) + exit_code: int = main(["--project-dir", str(tmp_path), "compile", "--no-cache", "--json"]) + remaining: list[str] = staging_directories(target) + + payload: dict[str, object] = json.loads(capsys.readouterr().out) + assert exit_code == 0 + assert payload["has_errors"] is False + assert tuple(remaining) == test_case.expected_remaining + assert len(list((target / "compiled" / "models").glob("*.sql"))) == test_case.model_count + + +@pytest.mark.skipif( + not has_second_filesystem(_SECOND_FILESYSTEM), + reason="needs /dev/shm on a different filesystem from the temporary directory", +) +@pytest.mark.parametrize( + "test_case", + (CrossDeviceArtifactsCompileTestCase("compiled directory on another filesystem", 128),), + ids=lambda case: case.description, +) +def test_given_compiled_dir_on_other_filesystem_when_compiling_then_artifacts_are_published( + test_case: CrossDeviceArtifactsCompileTestCase, + tmp_path: Path, + capsys: pytest.CaptureFixture[str], +) -> None: + write_prepared_artifacts_project(project_dir=tmp_path, model_count=test_case.model_count) + with TemporaryDirectory(dir=_SECOND_FILESYSTEM) as other_filesystem: + compiled: Path = Path(other_filesystem) / "compiled" + compiled.mkdir() + (tmp_path / "target").mkdir() + (tmp_path / "target" / "compiled").symlink_to(compiled, target_is_directory=True) + exit_code: int = main(["--project-dir", str(tmp_path), "compile", "--no-cache", "--json"]) + published: list[Path] = sorted((compiled / "models").glob("*.sql")) + leftovers: list[str] = sorted(path.name for path in compiled.rglob(".*")) + with duckdb.connect() as connection: + last_value: object = connection.execute( + published[-1].read_text(encoding="utf-8") + ).fetchone() + device_differs: bool = os.stat(compiled).st_dev != os.stat(tmp_path).st_dev + + payload: dict[str, object] = json.loads(capsys.readouterr().out) + assert exit_code == test_case.expected_exit_code + assert payload["has_errors"] is False + assert device_differs + assert (tmp_path / "target" / "compiled").is_symlink() + assert len(published) == test_case.model_count + assert last_value == (test_case.model_count - 1,) + assert leftovers == [] + assert staging_directories(tmp_path / "target") == [] + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__, "-n", "auto", "--dist", "loadfile"])) diff --git a/tests/unit/src/sqlbuild/cli/commands/main/compile/_test_types.py b/tests/unit/src/sqlbuild/cli/commands/main/compile/_test_types.py index efc2c81fc4..91edadb5c1 100644 --- a/tests/unit/src/sqlbuild/cli/commands/main/compile/_test_types.py +++ b/tests/unit/src/sqlbuild/cli/commands/main/compile/_test_types.py @@ -166,3 +166,26 @@ class RuleGatedTestPlanningTestCase: expected_codes: tuple[str, ...] expected_message_fragment: str expected_inline_planning_calls: int + + +@dataclass(frozen=True) +class CrossDeviceStagedPublishTestCase: + description: str + existing_files: dict[str, str] + expected_files: dict[str, str] + expected_removed: tuple[str, ...] + + +@dataclass(frozen=True) +class TamperedStagedPublishTestCase: + description: str + existing_files: dict[str, str] + removed_staged_files: tuple[str, ...] + expected_kept_files: dict[str, str] + + +@dataclass(frozen=True) +class StagingInterleavingTestCase: + description: str + existing_files: dict[str, str] + expected_files: dict[str, str] diff --git a/tests/unit/src/sqlbuild/cli/commands/main/compile/helpers.py b/tests/unit/src/sqlbuild/cli/commands/main/compile/helpers.py index 0ae7db5081..ec3f18b1b1 100644 --- a/tests/unit/src/sqlbuild/cli/commands/main/compile/helpers.py +++ b/tests/unit/src/sqlbuild/cli/commands/main/compile/helpers.py @@ -1,5 +1,8 @@ from __future__ import annotations +import errno +import os +from collections.abc import Callable from dataclasses import replace from pathlib import Path from typing import Any @@ -560,3 +563,35 @@ def prepare_rule_gated_compile_project(root: Path, *, test_header: str) -> Path: encoding="utf-8", ) return project_dir + + +def write_relative_files(*, root: Path, files: dict[str, str]) -> None: + """Write each relative path under root with its text contents.""" + + for relative, contents in files.items(): + path: Path = root / relative + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(contents, encoding="utf-8") + + +def cross_device_replace( + *, staging_root: Path, real_replace: Callable[[Any, Any], None] +) -> Callable[[Any, Any], None]: + """Return an os.replace that rejects moves out of staging_root like a second filesystem.""" + + staging_prefix: str = os.fspath(staging_root) + replacements: dict[bool, Callable[[Any, Any], None]] = { + True: cross_device_rename, + False: real_replace, + } + + def replace_file(source: Any, destination: Any) -> None: + replacements[os.fspath(source).startswith(staging_prefix)](source, destination) + + return replace_file + + +def cross_device_rename(source: Any, destination: Any) -> None: + """Reject a move as if source and destination were on different filesystems.""" + + raise OSError(errno.EXDEV, "Invalid cross-device link") diff --git a/tests/unit/src/sqlbuild/cli/commands/main/compile/test_staged_artifact_publication.py b/tests/unit/src/sqlbuild/cli/commands/main/compile/test_staged_artifact_publication.py new file mode 100644 index 0000000000..adffe372eb --- /dev/null +++ b/tests/unit/src/sqlbuild/cli/commands/main/compile/test_staged_artifact_publication.py @@ -0,0 +1,226 @@ +"""Staged cold-compile artifacts publish safely across filesystems and concurrent cleanup.""" + +from __future__ import annotations + +import os +import shutil +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import pytest +from filelock import FileLock, Timeout + +from sqlbuild.adapters.duckdb.classes.duckdb_adapter import DuckDbAdapter +from sqlbuild.cli.commands._helpers.compile.target_writer import ( + publish_static_compile_target, + staged_artifact_files, + write_static_compile_target, +) +from sqlbuild.cli.commands.classes import prepared_compile_artifacts +from sqlbuild.cli.commands.classes.prepared_compile_artifacts import PreparedCompileArtifacts +from sqlbuild.cli.commands.exceptions import StagedArtifactsChangedError +from sqlbuild.cli.output.models import WrittenTarget +from tests.unit.src.sqlbuild.cli.commands.main.compile._test_types import ( + CrossDeviceStagedPublishTestCase, + StagingInterleavingTestCase, + TamperedStagedPublishTestCase, +) +from tests.unit.src.sqlbuild.cli.commands.main.compile.helpers import ( + build_static_target_writer_project, + cross_device_rename, + cross_device_replace, + read_target_files, + write_relative_files, +) + +_MODEL_FILE: str = "compiled/models/staging/orders.sql" +_FUNCTION_FILE: str = "compiled/functions/sql/is_completed_order.sql" +_PUBLISHED_FILES: dict[str, str] = { + _MODEL_FILE: "SELECT 2 AS order_id\n", + _FUNCTION_FILE: ( + "CREATE OR REPLACE MACRO analytics.is_completed_order(order_status) AS (\n" + "order_status = 'completed'\n" + ")\n" + ), +} + + +@pytest.mark.parametrize( + "test_case", + [ + CrossDeviceStagedPublishTestCase( + description="new files join an existing tree on another filesystem", + existing_files={"compiled/models/obsolete.sql": "SELECT 0\n"}, + expected_files=_PUBLISHED_FILES, + expected_removed=("compiled/models/obsolete.sql",), + ), + CrossDeviceStagedPublishTestCase( + description="a cold tree that cannot be moved is copied file by file", + existing_files={}, + expected_files=_PUBLISHED_FILES, + expected_removed=(), + ), + ], + ids=lambda case: case.description, +) +def test_given_cross_device_target_when_publishing_staged_artifacts_then_files_are_copied( + test_case: CrossDeviceStagedPublishTestCase, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + staging: Path = tmp_path / "staging" + target: Path = tmp_path / "target" + prepared: WrittenTarget = write_static_compile_target( + target_dir=staging, adapter=DuckDbAdapter(), project=build_static_target_writer_project() + ) + expected_names: frozenset[str] = staged_artifact_files(target_dir=staging) + write_relative_files(root=target, files=test_case.existing_files) + monkeypatch.setattr( + os, "replace", cross_device_replace(staging_root=staging, real_replace=os.replace) + ) + monkeypatch.setattr(os, "rename", cross_device_rename) + + written: WrittenTarget = publish_static_compile_target( + prepared=prepared, target_dir=target, manifest=None, expected_files=expected_names + ) + + assert written.target_dir == target + assert read_target_files(target, test_case.expected_files) == test_case.expected_files + assert not any((target / relative).exists() for relative in test_case.expected_removed) + assert sorted(path.name for path in target.rglob("*.tmp")) == [] + + +@pytest.mark.parametrize( + "test_case", + [ + TamperedStagedPublishTestCase( + description="a staged file deleted after preparation", + existing_files={"compiled/models/obsolete.sql": "SELECT 0\n"}, + removed_staged_files=(_MODEL_FILE,), + expected_kept_files={"compiled/models/obsolete.sql": "SELECT 0\n"}, + ) + ], + ids=lambda case: case.description, +) +def test_given_changed_staging_tree_when_publishing_then_nothing_stale_is_deleted( + test_case: TamperedStagedPublishTestCase, tmp_path: Path +) -> None: + staging: Path = tmp_path / "staging" + target: Path = tmp_path / "target" + prepared: WrittenTarget = write_static_compile_target( + target_dir=staging, adapter=DuckDbAdapter(), project=build_static_target_writer_project() + ) + expected_names: frozenset[str] = staged_artifact_files(target_dir=staging) + write_relative_files(root=target, files=test_case.existing_files) + for relative in test_case.removed_staged_files: + (staging / relative).unlink() + + with pytest.raises(StagedArtifactsChangedError): + publish_static_compile_target( + prepared=prepared, target_dir=target, manifest=None, expected_files=expected_names + ) + + assert read_target_files(target, test_case.expected_kept_files) == ( + test_case.expected_kept_files + ) + assert not (target / _MODEL_FILE).exists() + + +@pytest.mark.parametrize( + "test_case", + [ + StagingInterleavingTestCase( + description="cleanup during staging creation keeps the live compile's tree", + existing_files={ + ".sqlbuild-staging-abandoned/compiled/models/orders.sql": "SELECT 0\n", + "compiled/models/obsolete.sql": "SELECT 0\n", + }, + expected_files=_PUBLISHED_FILES, + ) + ], + ids=lambda case: case.description, +) +def test_given_cleanup_interleaved_with_staging_when_compiling_then_live_staging_survives( + test_case: StagingInterleavingTestCase, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + target: Path = tmp_path / "target" + write_relative_files(root=target, files=test_case.existing_files) + survived: list[bool] = [] + create_staging: Callable[[Path], None] = prepared_compile_artifacts._create_staging_directory + remove_tree: Callable[..., None] = shutil.rmtree + + def create_then_clean(staging_dir: Path) -> None: + create_staging(staging_dir) + PreparedCompileArtifacts(enabled=True).remove_abandoned_staging(target_dir=target) + survived.append((staging_dir / ".owner").is_file()) + + def remove_while_locked(path: Any, *args: Any, **kwargs: Any) -> None: + with pytest.raises(Timeout): + FileLock(f"{os.fspath(path)}.lock", timeout=0).acquire() + remove_tree(path, *args, **kwargs) + + monkeypatch.setattr(prepared_compile_artifacts, "_MIN_PREPARED_ARTIFACT_MODELS", 1) + monkeypatch.setattr(prepared_compile_artifacts, "_create_staging_directory", create_then_clean) + monkeypatch.setattr(shutil, "rmtree", remove_while_locked) + + with PreparedCompileArtifacts(enabled=True) as artifacts: + artifacts.remove_abandoned_staging(target_dir=target) + artifacts.start( + project=build_static_target_writer_project(), + adapter=DuckDbAdapter(), + target_dir=target, + ) + written: WrittenTarget | None = artifacts.publish(target_dir=target, manifest=None) + + assert survived == [True] + assert written is not None + assert read_target_files(target, test_case.expected_files) == test_case.expected_files + assert not (target / "compiled" / "models" / "obsolete.sql").exists() + assert sorted(path.name for path in target.glob(".sqlbuild-staging-*")) == [] + + +@pytest.mark.parametrize( + "test_case", + [ + TamperedStagedPublishTestCase( + description="the owner marker removed from a live staging tree", + existing_files={"compiled/models/obsolete.sql": "SELECT 0\n"}, + removed_staged_files=(".owner",), + expected_kept_files={"compiled/models/obsolete.sql": "SELECT 0\n"}, + ) + ], + ids=lambda case: case.description, +) +def test_given_staging_owner_lost_when_publishing_then_compile_fails_without_deleting( + test_case: TamperedStagedPublishTestCase, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + target: Path = tmp_path / "target" + write_relative_files(root=target, files=test_case.existing_files) + monkeypatch.setattr(prepared_compile_artifacts, "_MIN_PREPARED_ARTIFACT_MODELS", 1) + + with PreparedCompileArtifacts(enabled=True) as artifacts: + artifacts.start( + project=build_static_target_writer_project(), + adapter=DuckDbAdapter(), + target_dir=target, + ) + _ = artifacts.planning_diagnostics() + for relative in test_case.removed_staged_files: + for staged in target.glob(f".sqlbuild-staging-*/{relative}"): + staged.unlink() + with pytest.raises(StagedArtifactsChangedError): + artifacts.publish(target_dir=target, manifest=None) + + assert read_target_files(target, test_case.expected_kept_files) == ( + test_case.expected_kept_files + ) + assert not (target / _MODEL_FILE).exists() + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-vv"])) From 30b8f5591cbba2ab115bcec6c062f563b6eae83e Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sat, 3 Oct 2026 04:26:07 +0100 Subject: [PATCH 2/4] perf(compile): scan SQL-test marker text ranges only when a marker needs them SQL-test planning rescanned each model's comments and quoted text once per marker kind even when that kind never occurs, and checked every match against every range. Scan lazily on the first candidate match, share one scan across the unresolved-reference checks, and look ranges up by binary search. Rendered test SQL is byte-identical. --- .../compiler/_helpers/sql_tests/markers.rs | 47 +++++++++--- .../compiler/_helpers/sql_tests/planning.rs | 10 +-- .../src/compiler/tests.rs | 2 + .../compiler/tests/test_sql_test_markers.rs | 72 +++++++++++++++++++ .../src/compiler/tests/test_types.rs | 13 ++++ 5 files changed, 130 insertions(+), 14 deletions(-) create mode 100644 crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_markers.rs diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/markers.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/markers.rs index de59e2f7ad..362a9f285b 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/markers.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/markers.rs @@ -19,7 +19,7 @@ pub(crate) fn replace_callable_markers( where F: FnMut(&str, &str) -> Option<(String, bool)>, { - let protected = protected_ranges(syntax, sql); + let mut protected = ProtectedRanges::new(syntax, sql); let mut output = String::with_capacity(sql.len()); let mut reached: HashSet = HashSet::new(); let mut cursor = 0; @@ -27,7 +27,7 @@ where let Some(full) = captures.get(0) else { continue; }; - if full.start() < cursor || in_protected_range(full.start(), &protected) { + if full.start() < cursor || protected.contains(full.start()) { continue; } let suffix_start = skip_whitespace(sql, full.end()); @@ -107,14 +107,14 @@ fn replace_marker_captures( where F: FnMut(&Captures<'_>) -> Option, { - let protected = protected_ranges(syntax, sql); + let mut protected = ProtectedRanges::new(syntax, sql); let mut output = String::with_capacity(sql.len()); let mut cursor = 0; for captures in pattern.captures_iter(sql) { let Some(full) = captures.get(0) else { continue; }; - if in_protected_range(full.start(), &protected) { + if protected.contains(full.start()) { continue; } let Some(value) = replacement(&captures) else { @@ -129,13 +129,18 @@ where } pub(crate) fn marker_names(pattern: &Regex, syntax: &LexicalSyntax, sql: &str) -> Vec { - let protected = protected_ranges(syntax, sql); + marker_names_in(pattern, &mut ProtectedRanges::new(syntax, sql)) +} + +/// Marker names outside the comments and quoted text of the SQL that `protected` scans. +pub(crate) fn marker_names_in(pattern: &Regex, protected: &mut ProtectedRanges<'_>) -> Vec { + let sql = protected.sql; let mut names: Vec = Vec::new(); for captures in pattern.captures_iter(sql) { let Some(full) = captures.get(0) else { continue; }; - if in_protected_range(full.start(), &protected) { + if protected.contains(full.start()) { continue; } if let Some(name) = captures.get(1) { @@ -149,8 +154,32 @@ pub(crate) fn protected_ranges(syntax: &LexicalSyntax, sql: &str) -> Vec<(usize, dialect_non_code_ranges(sql, syntax) } +/// Whether `index` falls inside one of the sorted, non-overlapping `ranges`. pub(crate) fn in_protected_range(index: usize, ranges: &[(usize, usize)]) -> bool { - ranges - .iter() - .any(|(start, end)| index >= *start && index < *end) + let following = ranges.partition_point(|(start, _)| *start <= index); + following > 0 && index < ranges[following - 1].1 +} + +/// Comment and quoted-text ranges of one SQL text, scanned only when a marker needs them. +pub(crate) struct ProtectedRanges<'a> { + sql: &'a str, + syntax: &'a LexicalSyntax, + ranges: Option>, +} + +impl<'a> ProtectedRanges<'a> { + pub(crate) fn new(syntax: &'a LexicalSyntax, sql: &'a str) -> Self { + Self { + sql, + syntax, + ranges: None, + } + } + + pub(crate) fn contains(&mut self, index: usize) -> bool { + let ranges = self + .ranges + .get_or_insert_with(|| protected_ranges(self.syntax, self.sql)); + in_protected_range(index, ranges) + } } diff --git a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/planning.rs b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/planning.rs index 7ca1946254..8e477285cb 100644 --- a/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/planning.rs +++ b/crates/sqlbuild-rules-native/src/compiler/_helpers/sql_tests/planning.rs @@ -19,7 +19,7 @@ use crate::compiler::_helpers::sql_tests::helper_scope::{ ScopeGraph, helper_scope_ctes, merged_scoped_ctes, }; use crate::compiler::_helpers::sql_tests::markers::{ - in_protected_range, marker_names, protected_ranges, replace_callable_markers, + ProtectedRanges, marker_names, marker_names_in, replace_callable_markers, replace_dbt_ref_markers, replace_named_markers, }; use crate::compiler::_helpers::sql_tests::rendering::{ @@ -1620,7 +1620,8 @@ fn unresolved_reference_warnings(request: UnresolvedReferenceRequest<'_>) -> Vec }); } }; - for name in marker_names(&patterns.reference, &patterns.lexical, sql) { + let mut protected = ProtectedRanges::new(&patterns.lexical, sql); + for name in marker_names_in(&patterns.reference, &mut protected) { let message = format!( "test '{test_name}': model '{model_name}' references __ref('{name}') which has no mock and is not in the expected chain" ); @@ -1630,19 +1631,18 @@ fn unresolved_reference_warnings(request: UnresolvedReferenceRequest<'_>) -> Vec (&patterns.source, SOURCE_FUNCTION), (&patterns.seed, SEED_FUNCTION), ] { - for name in marker_names(pattern, &patterns.lexical, sql) { + for name in marker_names_in(pattern, &mut protected) { let message = format!( "test '{test_name}': model '{model_name}' references {function_name}('{name}') which has no mock" ); warn(function_name, name, message); } } - let protected = protected_ranges(&patterns.lexical, sql); for captures in patterns.dbt_reference.captures_iter(sql) { let Some(full) = captures.get(0) else { continue; }; - if in_protected_range(full.start(), &protected) { + if protected.contains(full.start()) { continue; } let Some(first) = captures.get(1).map(|value| value.as_str()) else { diff --git a/crates/sqlbuild-rules-native/src/compiler/tests.rs b/crates/sqlbuild-rules-native/src/compiler/tests.rs index e5221df173..4888371b81 100644 --- a/crates/sqlbuild-rules-native/src/compiler/tests.rs +++ b/crates/sqlbuild-rules-native/src/compiler/tests.rs @@ -14,6 +14,8 @@ mod sql_test_cte_slices; mod sql_test_expected_columns; #[path = "tests/test_sql_test_extraction.rs"] mod sql_test_extraction; +#[path = "tests/test_sql_test_markers.rs"] +mod sql_test_markers; #[path = "tests/test_sql_test_planning.rs"] mod sql_test_planning; #[path = "tests/test_sql_test_rendering.rs"] diff --git a/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_markers.rs b/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_markers.rs new file mode 100644 index 0000000000..17f55cbbfc --- /dev/null +++ b/crates/sqlbuild-rules-native/src/compiler/tests/test_sql_test_markers.rs @@ -0,0 +1,72 @@ +use regex::Regex; + +use crate::compiler::_helpers::sql_tests::markers::{ + in_protected_range, marker_names, replace_named_markers, +}; +use crate::compiler::tests::test_types::{MarkerNamesTestCase, ProtectedRangeTestCase}; +use crate::sql_scan::models::LexicalSyntax; + +#[test] +fn given_markers_around_comments_and_strings_when_scanning_then_only_code_markers_are_used() { + let test_cases = [ + MarkerNamesTestCase { + description: "no markers leaves the SQL unchanged", + sql: "SELECT 'a' AS label -- note\nFROM orders", + expected_names: &[], + expected_replaced_sql: "SELECT 'a' AS label -- note\nFROM orders", + }, + MarkerNamesTestCase { + description: "markers between many comments and strings", + sql: "SELECT '__ref(\"quoted\")' AS a, /* __ref(\"hidden\") */ b\nFROM __ref(\"orders\") -- __ref(\"commented\")\nJOIN __ref(\"customers\") ON 'x' = 'y'\nWHERE c = '__ref(\"text\")' AND d IN (SELECT id FROM __ref(\"products\"))", + expected_names: &["orders", "customers", "products"], + expected_replaced_sql: "SELECT '__ref(\"quoted\")' AS a, /* __ref(\"hidden\") */ b\nFROM mock_orders -- __ref(\"commented\")\nJOIN mock_customers ON 'x' = 'y'\nWHERE c = '__ref(\"text\")' AND d IN (SELECT id FROM mock_products)", + }, + ]; + let pattern = Regex::new(r#"(?i)__ref\("([^"]+)"\)"#).expect("valid marker pattern"); + let syntax = LexicalSyntax { + line_comment_prefixes: vec!["--".to_string()], + ..LexicalSyntax::default() + }; + for test_case in test_cases { + assert_eq!( + marker_names(&pattern, &syntax, test_case.sql), + test_case.expected_names, + "{}", + test_case.description + ); + assert_eq!( + replace_named_markers(test_case.sql, &pattern, &syntax, |name| Some(format!( + "mock_{name}" + ))), + test_case.expected_replaced_sql, + "{}", + test_case.description + ); + } +} + +#[test] +fn given_sorted_ranges_when_checking_offsets_then_only_offsets_inside_a_range_are_protected() { + let test_cases = [ + ProtectedRangeTestCase { + description: "no ranges protect nothing", + ranges: &[], + expected_protected: &[], + }, + ProtectedRangeTestCase { + description: "starts are inclusive and ends exclusive across several ranges", + ranges: &[(2, 4), (6, 7), (10, 15)], + expected_protected: &[2, 3, 6, 10, 11, 12, 13, 14], + }, + ]; + for test_case in test_cases { + let protected: Vec = (0..17) + .filter(|index| in_protected_range(*index, test_case.ranges)) + .collect(); + assert_eq!( + protected, test_case.expected_protected, + "{}", + test_case.description + ); + } +} diff --git a/crates/sqlbuild-rules-native/src/compiler/tests/test_types.rs b/crates/sqlbuild-rules-native/src/compiler/tests/test_types.rs index ae64b2147e..3f96531f7c 100644 --- a/crates/sqlbuild-rules-native/src/compiler/tests/test_types.rs +++ b/crates/sqlbuild-rules-native/src/compiler/tests/test_types.rs @@ -60,3 +60,16 @@ pub(crate) struct CteRenameTestCase { pub(crate) renamed_index: usize, pub(crate) expected_sql: Option<&'static str>, } + +pub(crate) struct MarkerNamesTestCase { + pub(crate) description: &'static str, + pub(crate) sql: &'static str, + pub(crate) expected_names: &'static [&'static str], + pub(crate) expected_replaced_sql: &'static str, +} + +pub(crate) struct ProtectedRangeTestCase { + pub(crate) description: &'static str, + pub(crate) ranges: &'static [(usize, usize)], + pub(crate) expected_protected: &'static [usize], +} From b8c701a9d9e3e5680c0eb38db5d89f8935ff2372 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sat, 3 Oct 2026 04:37:17 +0100 Subject: [PATCH 3/4] perf(compile): pause cyclic garbage collection during the compile command A cold compile keeps millions of objects alive until it finishes, so every full cyclic collection rescans that heap without freeing it (about 2-4 s of GIL-held time on the 3,000- and 5,000-model dense benchmarks). Disable the cyclic collector for the duration of the compile command and restore the previous state afterwards; reference counting still frees acyclic garbage. --- .../commands/_helpers/compile/collection.py | 20 +++++++ .../cli/commands/main/project/_compile.py | 2 + .../cli/commands/main/compile/_test_types.py | 8 +++ .../main/compile/test_compile_command.py | 53 +++++++++++++++++++ 4 files changed, 83 insertions(+) create mode 100644 src/sqlbuild/cli/commands/_helpers/compile/collection.py diff --git a/src/sqlbuild/cli/commands/_helpers/compile/collection.py b/src/sqlbuild/cli/commands/_helpers/compile/collection.py new file mode 100644 index 0000000000..7bfe2bb3c7 --- /dev/null +++ b/src/sqlbuild/cli/commands/_helpers/compile/collection.py @@ -0,0 +1,20 @@ +"""Pause cyclic garbage collection while one compile builds its long-lived project state.""" + +from __future__ import annotations + +import gc +from collections.abc import Iterator +from contextlib import contextmanager + + +@contextmanager +def paused_cyclic_collection() -> Iterator[None]: + """Skip full-heap cycle scans of long-lived compile state, then restore the prior state.""" + + was_enabled: bool = gc.isenabled() + gc.disable() + try: + yield + finally: + if was_enabled: + gc.enable() diff --git a/src/sqlbuild/cli/commands/main/project/_compile.py b/src/sqlbuild/cli/commands/main/project/_compile.py index 666180c987..39c537022e 100644 --- a/src/sqlbuild/cli/commands/main/project/_compile.py +++ b/src/sqlbuild/cli/commands/main/project/_compile.py @@ -7,6 +7,7 @@ from dataclasses import replace from pathlib import Path +from sqlbuild.cli.commands._helpers.compile.collection import paused_cyclic_collection from sqlbuild.cli.commands._helpers.compile.output import ( format_compile_json, format_compile_text, @@ -49,6 +50,7 @@ def run_compile(request: CompileCommandRequest) -> int: ) try: with ( + paused_cyclic_collection(), collect_compile_timings() as detailed_timings, PreparedCompileArtifacts( enabled=not request.profile_flags.skip_write diff --git a/tests/unit/src/sqlbuild/cli/commands/main/compile/_test_types.py b/tests/unit/src/sqlbuild/cli/commands/main/compile/_test_types.py index 91edadb5c1..2dc1abf35d 100644 --- a/tests/unit/src/sqlbuild/cli/commands/main/compile/_test_types.py +++ b/tests/unit/src/sqlbuild/cli/commands/main/compile/_test_types.py @@ -189,3 +189,11 @@ class StagingInterleavingTestCase: description: str existing_files: dict[str, str] expected_files: dict[str, str] + + +@dataclass(frozen=True) +class CompileCyclicCollectionTestCase: + description: str + collection_enabled_before: bool + expected_enabled_during: bool + expected_enabled_after: bool diff --git a/tests/unit/src/sqlbuild/cli/commands/main/compile/test_compile_command.py b/tests/unit/src/sqlbuild/cli/commands/main/compile/test_compile_command.py index dd8663809e..01b7ed6c4a 100644 --- a/tests/unit/src/sqlbuild/cli/commands/main/compile/test_compile_command.py +++ b/tests/unit/src/sqlbuild/cli/commands/main/compile/test_compile_command.py @@ -1,5 +1,6 @@ from __future__ import annotations +import gc import json from collections.abc import Callable from pathlib import Path @@ -11,6 +12,7 @@ from sqlbuild.cli.commands._helpers.compile import pipeline as compile_pipeline from sqlbuild.cli.commands._helpers.compile import status as compile_status from sqlbuild.cli.commands.classes import prepared_compile_artifacts +from sqlbuild.cli.commands.main.project import _compile as compile_command from sqlbuild.cli.commands.types import CompileLineageMode from sqlbuild.cli.compile.main.run import run_compile from sqlbuild.cli.compile.models import ( @@ -22,6 +24,7 @@ from tests.unit.src.sqlbuild.cli.commands.main.compile._test_types import ( CompileColumnContractModeTestCase, CompileCommandTestCase, + CompileCyclicCollectionTestCase, CompileDagArtifactTestCase, CompileJsonDiagnosticsTestCase, CompileLineageModeTestCase, @@ -625,3 +628,53 @@ def counted_static_planning(**kwargs: Any) -> tuple[CompilerDiagnostic, ...]: assert payload["summary"]["errors"] == len(test_case.expected_codes) assert len(planning_calls) == test_case.expected_inline_planning_calls assert not (project_dir / "target" / "compiled" / "tests").exists() + + +@pytest.mark.parametrize( + "test_case", + [ + CompileCyclicCollectionTestCase( + description="enabled collection is paused and restored", + collection_enabled_before=True, + expected_enabled_during=False, + expected_enabled_after=True, + ), + CompileCyclicCollectionTestCase( + description="disabled collection stays disabled", + collection_enabled_before=False, + expected_enabled_during=False, + expected_enabled_after=False, + ), + ], + ids=lambda case: case.description, +) +def test_given_collection_state_when_compiling_then_cycles_are_paused_and_state_restored( + test_case: CompileCyclicCollectionTestCase, + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + project_dir: Path = prepare_static_compile_project(tmp_path) + enabled_during: list[bool] = [] + analyze: Callable[..., Any] = compile_pipeline.analyze_compile_project + + def observed_analyze(**kwargs: Any) -> Any: + enabled_during.append(gc.isenabled()) + return analyze(**kwargs) + + monkeypatch.setattr( + compile_pipeline, "resolve_adapter", lambda *args, **kwargs: NoConnectDuckDbAdapter() + ) + monkeypatch.setattr(compile_command, "analyze_compile_project", observed_analyze) + set_collection: dict[bool, Callable[[], None]] = {True: gc.enable, False: gc.disable} + was_enabled: bool = gc.isenabled() + set_collection[test_case.collection_enabled_before]() + try: + exit_code: int = run_compile(CompileCommandRequest(project_dir=project_dir)) + enabled_after: bool = gc.isenabled() + finally: + set_collection[was_enabled]() + + assert exit_code == 0 + assert enabled_during == [test_case.expected_enabled_during] + assert enabled_after is test_case.expected_enabled_after + assert (project_dir / "target" / "compiled" / "models" / "orders.sql").exists() From 63859ba4ff4d834b1b236e1e11297e3b44cc4116 Mon Sep 17 00:00:00 2001 From: Kevin Longe <273747162+kvlonge@users.noreply.github.com> Date: Sat, 3 Oct 2026 06:40:43 +0100 Subject: [PATCH 4/4] refactor(lint): share one line-start table helper across lint and discovery Native SQL lint, header lint and model discovery each kept a private copy of the line-start table and offset-to-position mapping; the native SQL lint copy scanned character by character in Python. Give discovery the single owner using str.find and reuse it from both lint modules. Positions are unchanged. --- .../discovery/_helpers/sql/model_files.py | 13 ++++++-- src/sqlbuild/lint/_helpers/native.py | 33 +++++-------------- src/sqlbuild/lint/_helpers/native_sql.py | 25 ++++---------- 3 files changed, 26 insertions(+), 45 deletions(-) diff --git a/src/sqlbuild/compiler/discovery/_helpers/sql/model_files.py b/src/sqlbuild/compiler/discovery/_helpers/sql/model_files.py index c429eef188..117965b523 100644 --- a/src/sqlbuild/compiler/discovery/_helpers/sql/model_files.py +++ b/src/sqlbuild/compiler/discovery/_helpers/sql/model_files.py @@ -333,7 +333,7 @@ def matched_model_output_column_locations( relative_path=relative_path, projection_ranges=projection_ranges, extract_implicit_alias_columns=extract_implicit_alias_columns, - line_starts=_line_starts(contents), + line_starts=source_line_starts(contents), ) @@ -487,7 +487,9 @@ def _location_for_absolute_span( ) -def _line_starts(contents: str) -> tuple[int, ...]: +def source_line_starts(contents: str) -> tuple[int, ...]: + """Return the offset where each line of contents starts.""" + starts: list[int] = [0] index: int = contents.find("\n") while index != -1: @@ -496,6 +498,13 @@ def _line_starts(contents: str) -> tuple[int, ...]: return tuple(starts) +def source_position(*, offset: int, line_starts: tuple[int, ...]) -> tuple[int, int]: + """Map a zero-based offset to a one-based line and column.""" + + line_index: int = bisect_right(line_starts, offset) - 1 + return line_index + 1, offset - line_starts[line_index] + 1 + + def _split_top_level_select_items(*, sql: str, start: int, end: int) -> tuple[tuple[int, int], ...]: items: list[tuple[int, int]] = [] depth: int = 0 diff --git a/src/sqlbuild/lint/_helpers/native.py b/src/sqlbuild/lint/_helpers/native.py index ed33796a3d..0187249572 100644 --- a/src/sqlbuild/lint/_helpers/native.py +++ b/src/sqlbuild/lint/_helpers/native.py @@ -4,7 +4,6 @@ import re import textwrap -from bisect import bisect_right from dataclasses import dataclass from pathlib import Path @@ -14,6 +13,8 @@ from sqlbuild.compiler.discovery._helpers.sql.model_files import ( parse_header_values, # noqa: FFL102 - the compiler owns the canonical native header grammar prepare_model_header_tokens, + source_line_starts, + source_position, ) from sqlbuild.compiler.discovery.exceptions import ModelSqlParseError from sqlbuild.lint.constants import ( @@ -87,7 +88,7 @@ def lint_native_headers( """Run all native header rules and return their violations.""" violations: list[LintViolation] = [] - line_starts: tuple[int, ...] = _line_starts(contents) + line_starts: tuple[int, ...] = source_line_starts(contents) header: HeaderSpan for header in headers: violations.extend( @@ -234,9 +235,7 @@ def _lint_header_values( try: values: dict[str, object] = _parse_header_values(kind=header.kind, header_text=header_text) except Exception as error: # noqa: BLE001 - any parse failure is a lint fault - position: tuple[int, int] = _offset_to_position( - offset=header.start, line_starts=line_starts - ) + position: tuple[int, int] = source_position(offset=header.start, line_starts=line_starts) return ( LintViolation( file_path=file_path, @@ -341,7 +340,7 @@ def _lint_header_whitespace( header_text: str = contents[header.start : header.end] if not any(line != line.rstrip() for line in _split_outside_quotes(text=header_text)): return () - position: tuple[int, int] = _offset_to_position(offset=header.start, line_starts=line_starts) + position: tuple[int, int] = source_position(offset=header.start, line_starts=line_starts) return ( LintViolation( file_path=file_path, @@ -592,8 +591,8 @@ def _relocate_leading_comment( ) faults: list[LintViolation] = [] if relocated_description.count("\n") + 1 > config.max_description_lines: - position: tuple[int, int] = _offset_to_position( - offset=header.start, line_starts=_line_starts(contents) + position: tuple[int, int] = source_position( + offset=header.start, line_starts=source_line_starts(contents) ) faults.append( LintViolation( @@ -699,8 +698,8 @@ def _violation_for_header_start( message: str, remediation: str, ) -> LintViolation: - position: tuple[int, int] = _offset_to_position( - offset=header.start, line_starts=_line_starts(contents) + position: tuple[int, int] = source_position( + offset=header.start, line_starts=source_line_starts(contents) ) return LintViolation( file_path=file_path, @@ -714,20 +713,6 @@ def _violation_for_header_start( ) -def _line_starts(contents: str) -> tuple[int, ...]: - starts: list[int] = [0] - index: int = contents.find("\n") - while index >= 0: - starts.append(index + 1) - index = contents.find("\n", index + 1) - return tuple(starts) - - -def _offset_to_position(*, offset: int, line_starts: tuple[int, ...]) -> tuple[int, int]: - line_index: int = bisect_right(line_starts, offset) - 1 - return line_index + 1, offset - line_starts[line_index] + 1 - - def _inner_header_text(*, header_text: str) -> str: open_index: int = header_text.find("(") close_index: int = header_text.rfind(")") diff --git a/src/sqlbuild/lint/_helpers/native_sql.py b/src/sqlbuild/lint/_helpers/native_sql.py index 6f27d2eb41..6867c5aa6e 100644 --- a/src/sqlbuild/lint/_helpers/native_sql.py +++ b/src/sqlbuild/lint/_helpers/native_sql.py @@ -3,7 +3,6 @@ from __future__ import annotations import re -from bisect import bisect_right from pathlib import Path from threading import Event from typing import Any, cast @@ -28,6 +27,7 @@ ) from sqlbuild.compiler.compile.main.map_expanded_offset import map_expanded_offset from sqlbuild.compiler.compile.models import MappedOffset +from sqlbuild.lint._helpers.native import source_line_starts, source_position from sqlbuild.lint._helpers.sqlbuild_tokens import interpolation_text_at from sqlbuild.lint.classes.stop_context import LintStopContext from sqlbuild.lint.constants import ( @@ -41,7 +41,6 @@ from sqlbuild.lint.models import LintBody, LintConfig, LintEdit, LintViolation _NATIVE_LINT_API_VERSION: int = 1 -_NEWLINE_CHARACTER: str = "\n" _UNUSED_CTE_CODE: str = "SQBRSQL005" _LONG_LITERAL_CODE: str = "SQBRSQL044" _PARSE_ERROR_POSITION_PATTERN: re.Pattern[str] = re.compile( @@ -130,7 +129,7 @@ def run_native_sql_lint( contents: str = contents_by_path[body.file_path] line_starts: tuple[int, ...] | None = line_starts_by_path.get(body.file_path) if line_starts is None: - line_starts = _line_starts(contents) + line_starts = source_line_starts(contents) line_starts_by_path[body.file_path] = line_starts if isinstance(response, NativeLintError): parse_violation: LintViolation | None = _parse_failure_violation( @@ -254,7 +253,7 @@ def _parse_failure_violation( expanded_offset += min(expanded_column - 1, len(expanded_lines[expanded_line - 1])) mapped: MappedOffset = map_expanded_offset(offset=expanded_offset, passes=body.passes) absolute_offset: int = body.body_start + mapped.offset - line, column = _offset_position(offset=absolute_offset, line_starts=line_starts) + line, column = source_position(offset=absolute_offset, line_starts=line_starts) return LintViolation( file_path=body.file_path, line=line, @@ -308,7 +307,7 @@ def _authored_violation( _SQLBUILD_HARNESS_CTE_PREFIXES, absolute_offset ): return None - line, column = _offset_position(offset=absolute_offset, line_starts=line_starts) + line, column = source_position(offset=absolute_offset, line_starts=line_starts) end_position: tuple[int, int] | None = _authored_end_position( start=start, end=end, @@ -422,14 +421,14 @@ def _authored_end_position( ) if token is None: return None - return _offset_position(offset=absolute_start + len(token), line_starts=line_starts) + return source_position(offset=absolute_start + len(token), line_starts=line_starts) if end <= start: return None mapped_last: MappedOffset = map_expanded_offset(offset=end - 1, passes=body.passes) if mapped_last.generated or mapped_last.offset - mapped_start.offset != end - start - 1: return None absolute_end: int = body.body_start + mapped_last.offset + 1 - return _offset_position(offset=absolute_end, line_starts=line_starts) + return source_position(offset=absolute_end, line_starts=line_starts) def _violation_message( @@ -441,15 +440,3 @@ def _violation_message( if token is None: return f"{message} {GENERATED_SQL_MESSAGE_SUFFIX}" return f"{message} {GENERATED_SQL_MESSAGE_TEMPLATE.format(token=token)}" - - -def _line_starts(value: str) -> tuple[int, ...]: - return ( - 0, - *(index + 1 for index, character in enumerate(value) if character == _NEWLINE_CHARACTER), - ) - - -def _offset_position(*, offset: int, line_starts: tuple[int, ...]) -> tuple[int, int]: - line_index: int = bisect_right(line_starts, offset) - 1 - return line_index + 1, offset - line_starts[line_index] + 1