|
19 | 19 | from types import FrameType |
20 | 20 | from typing import Any, NamedTuple, cast |
21 | 21 |
|
| 22 | +import duckdb |
22 | 23 | import pytest |
23 | 24 |
|
24 | 25 | from scripts.cold_compile_performance.main.read_compile_measurement import read_compile_measurement |
|
29 | 30 | from tests.e2e.src.sqlbuild.cli.commands.main.compile._test_types import ( |
30 | 31 | FreshProcessCompileCachePerformanceGuardTestCase, |
31 | 32 | SemanticCorpusCase, |
| 33 | + SetOperationModel, |
32 | 34 | ) |
33 | | -from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import prepare_inline_project |
| 35 | +from tests.e2e.src.sqlbuild.cli.commands.shared.helpers import prepare_inline_project, run_sqb |
34 | 36 |
|
35 | 37 | _DBT_SHAPED_SQL_SIZE_PROFILE: tuple[tuple[float, int], ...] = ( |
36 | 38 | (0.50, 1_800), |
@@ -2646,3 +2648,178 @@ def require_sql_analysis_output( |
2646 | 2648 | tuple((item["code"], item.get("path"), item.get("line")) for item in diagnostics), |
2647 | 2649 | text, |
2648 | 2650 | ) |
| 2651 | + |
| 2652 | + |
| 2653 | +_INDENTED_SET_OPERATION_PATTERN: re.Pattern[str] = re.compile( |
| 2654 | + r"^[ \t]+(?:UNION|EXCEPT|INTERSECT)\b.*$", re.MULTILINE |
| 2655 | +) |
| 2656 | + |
| 2657 | + |
| 2658 | +class SetOperationCompileResult(NamedTuple): |
| 2659 | + exit_code: int |
| 2660 | + diagnostics: tuple[tuple[str, str], ...] |
| 2661 | + column_counts: dict[str, int] |
| 2662 | + output: str |
| 2663 | + |
| 2664 | + |
| 2665 | +class SetOperationLifecycleResult(NamedTuple): |
| 2666 | + compiled: SetOperationCompileResult |
| 2667 | + build: subprocess.CompletedProcess[str] |
| 2668 | + relation_shapes: dict[str, tuple[int, int]] |
| 2669 | + rules: subprocess.CompletedProcess[str] |
| 2670 | + format: subprocess.CompletedProcess[str] |
| 2671 | + indented_set_operation_lines: tuple[str, ...] |
| 2672 | + recompiled: SetOperationCompileResult |
| 2673 | + |
| 2674 | + |
| 2675 | +def write_set_operation_project( |
| 2676 | + *, project_dir: Path, models: tuple[SetOperationModel, ...] |
| 2677 | +) -> Path: |
| 2678 | + """Write a DuckDB project with one view per model and return its database path.""" |
| 2679 | + |
| 2680 | + database: Path = project_dir / "orders.duckdb" |
| 2681 | + files: dict[str, str] = { |
| 2682 | + "sqlbuild_project.toml": ( |
| 2683 | + f'name = "orders"\nadapter = "duckdb"\n[connection]\ndatabase = "{database}"\n' |
| 2684 | + '[rules]\nselect = ["SQBRSQL020"]\n' |
| 2685 | + ), |
| 2686 | + } |
| 2687 | + for model in models: |
| 2688 | + files[f"models/{model.name}.sql"] = ( |
| 2689 | + f"MODEL (materialized view, schema analytics);\n\n{model.query_sql}\n" |
| 2690 | + ) |
| 2691 | + prepare_inline_project( |
| 2692 | + tmp_path=project_dir.parent, project_name=project_dir.name, repo_files=files |
| 2693 | + ) |
| 2694 | + return database |
| 2695 | + |
| 2696 | + |
| 2697 | +def compile_set_operation_project(*, project_dir: Path) -> SetOperationCompileResult: |
| 2698 | + """Compile without cache and return diagnostics and per-model output column counts.""" |
| 2699 | + |
| 2700 | + result: subprocess.CompletedProcess[str] = run_sqb( |
| 2701 | + project_dir=project_dir, command=("compile", "--no-cache", "--json") |
| 2702 | + ) |
| 2703 | + payload: dict[str, Any] = json.loads(result.stdout) |
| 2704 | + return SetOperationCompileResult( |
| 2705 | + exit_code=result.returncode, |
| 2706 | + diagnostics=tuple((item["code"], item["message"]) for item in payload["diagnostics"]), |
| 2707 | + column_counts={ |
| 2708 | + model["name"]: model["column_count"] for model in payload["resources"]["models"] |
| 2709 | + }, |
| 2710 | + output=result.stdout + result.stderr, |
| 2711 | + ) |
| 2712 | + |
| 2713 | + |
| 2714 | +def _indented_set_operation_lines(path: Path) -> list[str]: |
| 2715 | + lines: list[str] = _INDENTED_SET_OPERATION_PATTERN.findall(path.read_text(encoding="utf-8")) |
| 2716 | + return [f"{path.name}: {line}" for line in lines] |
| 2717 | + |
| 2718 | + |
| 2719 | +def run_set_operation_lifecycle( |
| 2720 | + *, project_dir: Path, models: tuple[SetOperationModel, ...] |
| 2721 | +) -> SetOperationLifecycleResult: |
| 2722 | + """Compile, build, run Rules, format, and recompile a set-operation project.""" |
| 2723 | + |
| 2724 | + database: Path = write_set_operation_project(project_dir=project_dir, models=models) |
| 2725 | + compiled: SetOperationCompileResult = compile_set_operation_project(project_dir=project_dir) |
| 2726 | + build: subprocess.CompletedProcess[str] = run_sqb(project_dir=project_dir, command=("build",)) |
| 2727 | + relation_shapes: dict[str, tuple[int, int]] = {} |
| 2728 | + with duckdb.connect(str(database), read_only=True) as connection: |
| 2729 | + for model in models: |
| 2730 | + relation: duckdb.DuckDBPyRelation = connection.sql( |
| 2731 | + f"SELECT * FROM analytics.{model.name}" |
| 2732 | + ) |
| 2733 | + relation_shapes[model.name] = (len(relation.columns), len(relation.fetchall())) |
| 2734 | + rules: subprocess.CompletedProcess[str] = run_sqb( |
| 2735 | + project_dir=project_dir, command=("rules", "run", "SQBRSQL020") |
| 2736 | + ) |
| 2737 | + models_dir: Path = project_dir / "models" |
| 2738 | + formatted: subprocess.CompletedProcess[str] = run_sqb( |
| 2739 | + project_dir=project_dir, command=("format", str(models_dir)) |
| 2740 | + ) |
| 2741 | + indented_set_operation_lines: list[str] = [] |
| 2742 | + for path in sorted(models_dir.glob("*.sql")): |
| 2743 | + indented_set_operation_lines.extend(_indented_set_operation_lines(path)) |
| 2744 | + return SetOperationLifecycleResult( |
| 2745 | + compiled=compiled, |
| 2746 | + build=build, |
| 2747 | + relation_shapes=relation_shapes, |
| 2748 | + rules=rules, |
| 2749 | + format=formatted, |
| 2750 | + indented_set_operation_lines=tuple(indented_set_operation_lines), |
| 2751 | + recompiled=compile_set_operation_project(project_dir=project_dir), |
| 2752 | + ) |
| 2753 | + |
| 2754 | + |
| 2755 | +class UnionFixtureCompileMeasurement(NamedTuple): |
| 2756 | + analysis_native_ms: int |
| 2757 | + sql_tests: int |
| 2758 | + errors: int |
| 2759 | + |
| 2760 | + |
| 2761 | +def _union_all_fixture_rows(*, row_count: int, status: str) -> str: |
| 2762 | + return " UNION ALL\n".join( |
| 2763 | + f" SELECT {row} AS order_id, CAST({row} AS DOUBLE) AS amount, '{status}' AS status" |
| 2764 | + for row in range(row_count) |
| 2765 | + ) |
| 2766 | + |
| 2767 | + |
| 2768 | +def write_union_fixture_test_project( |
| 2769 | + *, project_dir: Path, sql_test_count: int, fixture_row_count: int |
| 2770 | +) -> None: |
| 2771 | + """Write a DuckDB project whose SQL tests use long UNION ALL fixture chains.""" |
| 2772 | + |
| 2773 | + files: dict[str, str] = { |
| 2774 | + "sqlbuild_project.toml": ( |
| 2775 | + f'name = "orders"\nadapter = "duckdb"\n[connection]\n' |
| 2776 | + f'database = "{project_dir / "orders.duckdb"}"\n[rules]\nselect = []\n' |
| 2777 | + ), |
| 2778 | + "models/base_orders.sql": ( |
| 2779 | + "MODEL (materialized view);\n\n" |
| 2780 | + "SELECT 1 AS order_id, CAST(1 AS DOUBLE) AS amount, 'open' AS status\n" |
| 2781 | + ), |
| 2782 | + "models/orders.sql": ( |
| 2783 | + "MODEL (materialized view);\n\n" |
| 2784 | + 'SELECT b.order_id, b.amount, b.status FROM __ref("base_orders") AS b\n' |
| 2785 | + ), |
| 2786 | + } |
| 2787 | + for index in range(sql_test_count): |
| 2788 | + fixture_rows: str = _union_all_fixture_rows( |
| 2789 | + row_count=fixture_row_count, status=f"status_{index:02d}" |
| 2790 | + ) |
| 2791 | + files[f"tests/unit/orders_case_{index:02d}.sql"] = ( |
| 2792 | + f'TEST (name "orders_case_{index:02d}");\n\n' |
| 2793 | + f"WITH\n__ref__base_orders AS (\n{fixture_rows}\n),\n" |
| 2794 | + f"__expected__orders AS (\n{fixture_rows}\n)\nSELECT 1\n" |
| 2795 | + ) |
| 2796 | + prepare_inline_project( |
| 2797 | + tmp_path=project_dir.parent, project_name=project_dir.name, repo_files=files |
| 2798 | + ) |
| 2799 | + |
| 2800 | + |
| 2801 | +def measure_union_fixture_compile( |
| 2802 | + *, project_dir: Path, runs: int |
| 2803 | +) -> UnionFixtureCompileMeasurement: |
| 2804 | + """Return the fastest native analysis time of fresh uncached compiles.""" |
| 2805 | + |
| 2806 | + skip_actions: dict[bool, Callable[[], None]] = { |
| 2807 | + False: _continue_compile_benchmark, |
| 2808 | + True: _skip_compile_benchmark, |
| 2809 | + } |
| 2810 | + skip_actions[os.environ.get("SQLBUILD_SKIP_PERFORMANCE_TESTS") == "1"]() |
| 2811 | + measurements: list[UnionFixtureCompileMeasurement] = [] |
| 2812 | + for _ in range(runs): |
| 2813 | + result: subprocess.CompletedProcess[str] = run_installed_sqb( |
| 2814 | + project_dir=project_dir, args=("compile", "--no-cache", "--json"), env={} |
| 2815 | + ) |
| 2816 | + assert result.returncode == 0, result.stdout + result.stderr |
| 2817 | + payload: dict[str, Any] = json.loads(result.stdout) |
| 2818 | + measurements.append( |
| 2819 | + UnionFixtureCompileMeasurement( |
| 2820 | + analysis_native_ms=payload["compile_timings"]["analysis_native_ms"], |
| 2821 | + sql_tests=payload["summary"]["tests"], |
| 2822 | + errors=payload["summary"]["errors"], |
| 2823 | + ) |
| 2824 | + ) |
| 2825 | + return min(measurements) |
0 commit comments