perf(compile): plan SQL tests in the background in every compile mode

SQL-test planning only ran in the background for uncached projects of
128 to 5000 models with rules selected, so cache-on and large compiles
planned on the critical path in the write phase. Prepared artifacts now
fall back to planning SQL tests alone in the background, against the
published target and its artifact cache, and the write phase consumes
that plan at the point it would otherwise plan inline. Results and
errors are identical to inline planning.

Author: kvlonge

src/sqlbuild/cli/commands/_helpers/compile/pipeline.py

@@ -201,13 +201,12 @@ def _analyze_compile_project(
     )
     if not graph_failed:
         rules_config: RulesConfig = load_rules_config(project_dir=project_dir)
-        if (
-            prepared_artifacts is not None
-            and rules_config.select
-            and not any(item.is_error for item in core_diagnostics)
-        ):
+        if prepared_artifacts is not None and not any(item.is_error for item in core_diagnostics):
             prepared_artifacts.start(
-                project=graph.project, adapter=adapter, target_dir=project_dir / "target"
+                project=graph.project,
+                adapter=adapter,
+                target_dir=project_dir / "target",
+                stage_artifacts=bool(rules_config.select),
             )
         _ = start_compile_phase(status=status, message="Evaluating built-in and custom rules...")
         rules_start: float = time.monotonic()

src/sqlbuild/cli/commands/_helpers/compile/target_writer.py

@@ -6,6 +6,7 @@
 import json
 import os
 import uuid
+from collections.abc import Callable
 from pathlib import Path
 
 from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter
@@ -20,6 +21,8 @@
 )
 from sqlbuild.cli.commands.exceptions import StagedArtifactsChangedError
 from sqlbuild.cli.compile.models import (
+    PendingStaticSqlTest,
+    PlannedStaticSqlTests,
     SqlTestArtifactCacheRecord,
     SqlTestArtifactIdentityContext,
 )
@@ -124,8 +127,9 @@ def write_static_compile_target(
     adapter: BaseAdapter,
     project: CompiledProject,
     manifest: dict[str, object] | None = None,
+    planned_tests: Callable[[], PlannedStaticSqlTests] | None = None,
 ) -> WrittenTarget:
-    """Write offline compiled output files under target_dir."""
+    """Write offline compiled output files under target_dir, reusing SQL tests planned ahead."""
 
     remove_stale_files: bool = (target_dir / _COMPILED_DIR).is_dir()
     target_dir.mkdir(parents=True, exist_ok=True)
@@ -152,6 +156,7 @@ def write_static_compile_target(
         adapter=adapter,
         project=project,
         check_existing=remove_stale_files,
+        planned_tests=planned_tests,
     )
     managed_paths.update(test_paths)
     if remove_stale_files:
@@ -442,27 +447,60 @@ def _write_tests(
 
 
 def static_sql_test_planning_diagnostics(
-    *, target_dir: Path, adapter: BaseAdapter, project: CompiledProject
+    *,
+    target_dir: Path,
+    adapter: BaseAdapter,
+    project: CompiledProject,
+    planned_tests: Callable[[], PlannedStaticSqlTests] | None = None,
 ) -> tuple[CompilerDiagnostic, ...]:
     """Plan uncached SQL tests and report their errors without writing artifacts or cache."""
 
-    pending, _, _ = _partition_cached_static_tests(
-        tests_root=target_dir / _COMPILED_DIR / _TESTS_DIR, adapter=adapter, project=project
-    )
-    tests: tuple[CompiledSqlTest, ...] = tuple(test for test, _, _ in pending)
-    artifacts: tuple[NativeSqlTestArtifact, ...] = _plan_static_test_artifacts(
-        adapter=adapter, project=project, tests=tests
+    planned: PlannedStaticSqlTests = _resolve_planned_static_tests(
+        target_dir=target_dir, adapter=adapter, project=project, planned_tests=planned_tests
     )
     diagnostics: list[CompilerDiagnostic] = []
-    for test, artifact in zip(tests, artifacts, strict=True):
-        diagnostics.extend(_sql_test_artifact_diagnostics(test=test, artifact=artifact))
+    for pending, artifact in zip(planned.pending, planned.artifacts, strict=True):
+        diagnostics.extend(_sql_test_artifact_diagnostics(test=pending.test, artifact=artifact))
     return tuple(diagnostics)
 
 
+def plan_static_sql_tests(
+    *, target_dir: Path, adapter: BaseAdapter, project: CompiledProject
+) -> PlannedStaticSqlTests:
+    """Plan every SQL test whose cached artifact under target_dir cannot be reused."""
+
+    tests_root: Path = target_dir / _COMPILED_DIR / _TESTS_DIR
+    pending, cached_paths, cached_records = _partition_cached_static_tests(
+        tests_root=tests_root, adapter=adapter, project=project
+    )
+    return PlannedStaticSqlTests(
+        tests_root=tests_root,
+        pending=pending,
+        artifacts=_plan_static_test_artifacts(
+            adapter=adapter, project=project, tests=tuple(item.test for item in pending)
+        ),
+        cached_paths=frozenset(cached_paths),
+        cached_records=cached_records,
+    )
+
+
+def _resolve_planned_static_tests(
+    *,
+    target_dir: Path,
+    adapter: BaseAdapter,
+    project: CompiledProject,
+    planned_tests: Callable[[], PlannedStaticSqlTests] | None,
+) -> PlannedStaticSqlTests:
+    planned: PlannedStaticSqlTests | None = planned_tests() if planned_tests is not None else None
+    if planned is not None and planned.tests_root == target_dir / _COMPILED_DIR / _TESTS_DIR:
+        return planned
+    return plan_static_sql_tests(target_dir=target_dir, adapter=adapter, project=project)
+
+
 def _partition_cached_static_tests(
     *, tests_root: Path, adapter: BaseAdapter, project: CompiledProject
 ) -> tuple[
-    list[tuple[CompiledSqlTest, str | None, str | None]],
+    tuple[PendingStaticSqlTest, ...],
     set[Path],
     dict[str, SqlTestArtifactCacheRecord],
 ]:
@@ -484,7 +522,7 @@ def _partition_cached_static_tests(
             project=project,
             tests=project.sql_tests,
         )
-    pending: list[tuple[CompiledSqlTest, str | None, str | None]] = []
+    pending: list[PendingStaticSqlTest] = []
     for test in project.sql_tests:
         record_key: str | None = None
         artifact_identity: str | None = None
@@ -506,8 +544,12 @@ def _partition_cached_static_tests(
                     managed_paths.add(cached_path)
                     current_records[record_key] = cached_record
                     continue
-        pending.append((test, record_key, artifact_identity))
-    return pending, managed_paths, current_records
+        pending.append(
+            PendingStaticSqlTest(
+                test=test, record_key=record_key, artifact_identity=artifact_identity
+            )
+        )
+    return tuple(pending), managed_paths, current_records
 
 
 def _plan_static_test_artifacts(
@@ -534,20 +576,21 @@ def _write_static_tests(
     adapter: BaseAdapter,
     project: CompiledProject,
     check_existing: bool,
+    planned_tests: Callable[[], PlannedStaticSqlTests] | None,
 ) -> tuple[set[Path], tuple[CompilerDiagnostic, ...]]:
     """Write offline SQL-native test SQL and report uncached planning errors."""
 
     diagnostics: list[CompilerDiagnostic] = []
-    tests_root: Path = target_dir / _COMPILED_DIR / _TESTS_DIR
-    pending, managed_paths, current_records = _partition_cached_static_tests(
-        tests_root=tests_root, adapter=adapter, project=project
-    )
-    native_artifacts: tuple[NativeSqlTestArtifact, ...] = _plan_static_test_artifacts(
-        adapter=adapter, project=project, tests=tuple(test for test, _, _ in pending)
+    planned: PlannedStaticSqlTests = _resolve_planned_static_tests(
+        target_dir=target_dir, adapter=adapter, project=project, planned_tests=planned_tests
     )
-    for (test, record_key, artifact_identity), artifact in zip(
-        pending, native_artifacts, strict=True
-    ):
+    tests_root: Path = planned.tests_root
+    managed_paths: set[Path] = set(planned.cached_paths)
+    current_records: dict[str, SqlTestArtifactCacheRecord] = dict(planned.cached_records)
+    for pending, artifact in zip(planned.pending, planned.artifacts, strict=True):
+        test: CompiledSqlTest = pending.test
+        record_key: str | None = pending.record_key
+        artifact_identity: str | None = pending.artifact_identity
         test_path: Path = tests_root / compiled_sql_test_output_path(
             test=test,
             model_names=artifact.model_names,

src/sqlbuild/cli/commands/classes/prepared_compile_artifacts.py

@@ -1,4 +1,4 @@
-"""Own staged cold compile artifacts and publish them only at the normal write phase."""
+"""Prepare compile artifacts in the background and publish them only at the normal write phase."""
 
 from __future__ import annotations
 
@@ -14,11 +14,14 @@
 
 from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter
 from sqlbuild.cli.commands._helpers.compile.target_writer import (
+    plan_static_sql_tests,
     publish_static_compile_target,
     staged_artifact_files,
+    static_sql_test_planning_diagnostics,
     write_static_compile_target,
 )
 from sqlbuild.cli.commands.exceptions import StagedArtifactsChangedError
+from sqlbuild.cli.compile.models import PlannedStaticSqlTests
 from sqlbuild.cli.output.models import WrittenTarget
 from sqlbuild.compiler.compile.models import CompiledProject, CompilerDiagnostic
 
@@ -32,7 +35,7 @@
 
 
 class PreparedCompileArtifacts:
-    """Own one bounded preparation task and its disposable output directory."""
+    """Own one background preparation task: staged artifacts, or planned SQL tests only."""
 
     def __init__(self, *, enabled: bool) -> None:
         self.enabled: bool = enabled
@@ -41,6 +44,8 @@ def __init__(self, *, enabled: bool) -> None:
         self._created_target_dir: Path | None = None
         self._executor: ThreadPoolExecutor | None = None
         self._future: Future[_PreparedResult] | None = None
+        self._planned_tests: Future[PlannedStaticSqlTests] | None = None
+        self._planning_inputs: tuple[CompiledProject, BaseAdapter, Path] | None = None
 
     def __enter__(self) -> PreparedCompileArtifacts:
         return self
@@ -73,16 +78,48 @@ def remove_abandoned_staging(self, *, target_dir: Path) -> None:
             if not _remove_abandoned_directory(directory):
                 return
 
-    def start(self, *, project: CompiledProject, adapter: BaseAdapter, target_dir: Path) -> None:
-        """Stage artifacts beside target_dir so publication can move them instead of copying."""
+    def start(
+        self,
+        *,
+        project: CompiledProject,
+        adapter: BaseAdapter,
+        target_dir: Path,
+        stage_artifacts: bool,
+    ) -> None:
+        """Stage artifacts beside target_dir when worthwhile, otherwise plan SQL tests ahead."""
 
+        if not self.enabled:
+            return
         if (
-            not self.enabled
-            or project.compile_cache_dir is not None
-            or len(project.models) < _MIN_PREPARED_ARTIFACT_MODELS
-            or len(project.models) > _MAX_PREPARED_ARTIFACT_MODELS
+            stage_artifacts
+            and _stages_artifacts(project)
+            and self._start_staging(project=project, adapter=adapter, target_dir=target_dir)
         ):
             return
+        if project.sql_tests:
+            self._start_test_planning(project=project, adapter=adapter, target_dir=target_dir)
+
+    def _start_test_planning(
+        self, *, project: CompiledProject, adapter: BaseAdapter, target_dir: Path
+    ) -> None:
+        """Plan SQL tests against the published target while the compile continues."""
+
+        self._planning_inputs = (project, adapter, target_dir)
+        self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-tests")
+        context: Context = copy_context()
+
+        def plan() -> PlannedStaticSqlTests:
+            return context.run(
+                plan_static_sql_tests, target_dir=target_dir, adapter=adapter, project=project
+            )
+
+        self._planned_tests = self._executor.submit(plan)
+
+    def _start_staging(
+        self, *, project: CompiledProject, adapter: BaseAdapter, target_dir: Path
+    ) -> bool:
+        """Stage artifacts beside target_dir so publication can move them instead of copying."""
+
         staging_dir: Path = target_dir / f"{_STAGING_PREFIX}{uuid.uuid4().hex}"
         try:
             if not target_dir.is_dir():
@@ -92,7 +129,7 @@ def start(self, *, project: CompiledProject, adapter: BaseAdapter, target_dir: P
             self._owner_lock.acquire()
             _create_staging_directory(staging_dir)
         except OSError:
-            return
+            return False
         self._staging_dir = staging_dir
         self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="sqlbuild-artifacts")
         context: Context = copy_context()
@@ -107,10 +144,20 @@ def prepare() -> _PreparedResult:
             return written, staged_artifact_files(target_dir=staging_dir)
 
         self._future = self._executor.submit(prepare)
+        return True
 
     def publish(
         self, *, target_dir: Path, manifest: dict[str, object] | None
     ) -> WrittenTarget | None:
+        if self._planned_tests is not None and self._planning_inputs is not None:
+            project, adapter, _ = self._planning_inputs
+            return write_static_compile_target(
+                target_dir=target_dir,
+                adapter=adapter,
+                project=project,
+                manifest=manifest,
+                planned_tests=self._planned_tests.result,
+            )
         if self._future is None or self._staging_dir is None:
             return None
         try:
@@ -131,8 +178,16 @@ def publish(
         )
 
     def planning_diagnostics(self) -> tuple[CompilerDiagnostic, ...] | None:
-        """Return staged SQL test planning diagnostics without publishing any artifact."""
+        """Return background SQL test planning diagnostics without publishing any artifact."""
 
+        if self._planned_tests is not None and self._planning_inputs is not None:
+            project, adapter, target_dir = self._planning_inputs
+            return static_sql_test_planning_diagnostics(
+                target_dir=target_dir,
+                adapter=adapter,
+                project=project,
+                planned_tests=self._planned_tests.result,
+            )
         if self._future is None:
             return None
         try:
@@ -141,6 +196,15 @@ def planning_diagnostics(self) -> tuple[CompilerDiagnostic, ...] | None:
             return None
 
 
+def _stages_artifacts(project: CompiledProject) -> bool:
+    """Staging pays off for uncached, mid-sized projects whose memory stays within budget."""
+
+    return (
+        project.compile_cache_dir is None
+        and _MIN_PREPARED_ARTIFACT_MODELS <= len(project.models) <= _MAX_PREPARED_ARTIFACT_MODELS
+    )
+
+
 def _owner_lock_path(staging_dir: Path) -> Path:
     return staging_dir.with_name(f"{staging_dir.name}{_LOCK_SUFFIX}")
 

src/sqlbuild/cli/compile/models.py

@@ -8,10 +8,15 @@
 from sqlbuild.adapter.contract.classes.base_adapter import BaseAdapter
 from sqlbuild.cli.commands.types import CompileLineageMode
 from sqlbuild.cli.output.models import WrittenTarget
-from sqlbuild.compiler.compile.models import CompiledObjectKey, CompilerDiagnostic
+from sqlbuild.compiler.compile.models import (
+    CompiledObjectKey,
+    CompiledSqlTest,
+    CompilerDiagnostic,
+)
 from sqlbuild.compiler.discovery.models import DiscoveredProjectInputs
 from sqlbuild.compiler.lineage.models import ProjectColumnLineage
 from sqlbuild.compiler.pipeline.models import ProjectGraph
+from sqlbuild.compiler.planner.models import NativeSqlTestArtifact
 
 
 @dataclass(frozen=True)
@@ -90,3 +95,23 @@ class SqlTestArtifactIdentityContext:
 
     common_identity: str
     model_identities: dict[str, str]
+
+
+@dataclass(frozen=True)
+class PendingStaticSqlTest:
+    """One SQL test that needs a freshly planned artifact, with its optional cache identity."""
+
+    test: CompiledSqlTest
+    record_key: str | None
+    artifact_identity: str | None
+
+
+@dataclass(frozen=True)
+class PlannedStaticSqlTests:
+    """SQL-test artifacts planned for one target directory before any file is written."""
+
+    tests_root: Path
+    pending: tuple[PendingStaticSqlTest, ...]
+    artifacts: tuple[NativeSqlTestArtifact, ...]
+    cached_paths: frozenset[Path]
+    cached_records: dict[str, SqlTestArtifactCacheRecord]