|
1 | 1 | from __future__ import annotations |
2 | 2 |
|
3 | 3 | from dataclasses import replace |
| 4 | +from typing import Any |
4 | 5 |
|
5 | 6 | import pytest |
6 | 7 |
|
|
11 | 12 | CompiledProject, |
12 | 13 | ) |
13 | 14 | from sqlbuild.compiler.compile.types import CompiledResourceType |
| 15 | +from sqlbuild.compiler.lineage._helpers import rich_columns |
| 16 | +from sqlbuild.compiler.lineage._helpers.columns import _build_schema_mapping |
14 | 17 | from sqlbuild.compiler.lineage.main.columns import build_project_column_lineage |
15 | 18 | from sqlbuild.compiler.lineage.models import ColumnLineage, ModelColumnLineage, ProjectColumnLineage |
16 | 19 | from sqlbuild.compiler.lineage.types import ( |
17 | 20 | ColumnLineageConfidence, |
18 | 21 | ColumnLineageMode, |
19 | 22 | ColumnTransformKind, |
20 | 23 | ) |
| 24 | +from sqlbuild.compiler.sql_analysis.main.import_polyglot_sql import import_polyglot_sql |
21 | 25 | from tests.unit.src.sqlbuild.compiler.lineage._test_types import ( |
22 | 26 | ColumnLineageAnalyzerTestCase, |
23 | 27 | ExpectedCountTestCase, |
24 | 28 | ProjectLineageGraphTestCase, |
| 29 | + RichLineageAnalysisScalingTestCase, |
25 | 30 | SqlAnalysisDisabledLineageTestCase, |
26 | 31 | ) |
27 | 32 | from tests.unit.src.sqlbuild.compiler.lineage.helpers import ( |
28 | 33 | edge_label, |
| 34 | + make_chain_models, |
29 | 35 | make_compiled_model, |
30 | 36 | make_compiled_project, |
31 | 37 | make_compiled_seed, |
@@ -653,3 +659,78 @@ def test_given_sql_analysis_disabled_when_building_column_lineage_then_returns_n |
653 | 659 | result: ProjectColumnLineage | None = build_project_column_lineage(project=project) |
654 | 660 |
|
655 | 661 | assert (result is None) is test_case.expected_result_is_none |
| 662 | + |
| 663 | + |
| 664 | +@pytest.mark.parametrize( |
| 665 | + "test_case", |
| 666 | + [ |
| 667 | + RichLineageAnalysisScalingTestCase( |
| 668 | + description="builds the schema once and analyses each traced model against its refs", |
| 669 | + chain_length=4, |
| 670 | + unrelated_model_counts=(0, 250), |
| 671 | + expected_schema_builds=1, |
| 672 | + expected_analysis_calls=4, |
| 673 | + expected_schema_tables_per_call=( |
| 674 | + (), |
| 675 | + ("__sqlbuild_model__orders_0",), |
| 676 | + ("__sqlbuild_model__orders_1",), |
| 677 | + ("__sqlbuild_model__orders_2",), |
| 678 | + ), |
| 679 | + ) |
| 680 | + ], |
| 681 | + ids=lambda case: case.description, |
| 682 | +) |
| 683 | +def test_given_unrelated_models_when_tracing_rich_lineage_then_analysis_work_stays_linear( |
| 684 | + test_case: RichLineageAnalysisScalingTestCase, |
| 685 | + monkeypatch: pytest.MonkeyPatch, |
| 686 | +) -> None: |
| 687 | + polyglot_module: Any = import_polyglot_sql() |
| 688 | + analyze_query: Any = polyglot_module.analyze_query |
| 689 | + schema_builds: list[CompiledProject] = [] |
| 690 | + analysed_tables: list[tuple[str, ...]] = [] |
| 691 | + |
| 692 | + def counting_schema_mapping(project: CompiledProject) -> dict[str, dict[str, str]]: |
| 693 | + schema_builds.append(project) |
| 694 | + return _build_schema_mapping(project) |
| 695 | + |
| 696 | + def recording_analyze_query(sql: str, options: dict[str, Any]) -> object: |
| 697 | + analysed_tables.append(tuple(table["name"] for table in options["schema"]["tables"])) |
| 698 | + return analyze_query(sql, options) |
| 699 | + |
| 700 | + monkeypatch.setattr(rich_columns, "_build_schema_mapping", counting_schema_mapping) |
| 701 | + monkeypatch.setattr(polyglot_module, "analyze_query", recording_analyze_query) |
| 702 | + chain_names: frozenset[str] = frozenset( |
| 703 | + f"orders_{index}" for index in range(test_case.chain_length) |
| 704 | + ) |
| 705 | + edges_by_project_size: list[tuple[str, ...]] = [] |
| 706 | + for unrelated_model_count in test_case.unrelated_model_counts: |
| 707 | + schema_builds.clear() |
| 708 | + analysed_tables.clear() |
| 709 | + project: CompiledProject = make_compiled_project( |
| 710 | + models=make_chain_models( |
| 711 | + chain_length=test_case.chain_length, |
| 712 | + unrelated_model_count=unrelated_model_count, |
| 713 | + ) |
| 714 | + ) |
| 715 | + |
| 716 | + result: ProjectColumnLineage | None = build_project_column_lineage( |
| 717 | + project=project, model_names=chain_names |
| 718 | + ) |
| 719 | + |
| 720 | + assert result is not None |
| 721 | + assert len(schema_builds) == test_case.expected_schema_builds |
| 722 | + assert len(analysed_tables) == test_case.expected_analysis_calls |
| 723 | + assert tuple(analysed_tables) == test_case.expected_schema_tables_per_call |
| 724 | + edges_by_project_size.append( |
| 725 | + tuple( |
| 726 | + edge_label( |
| 727 | + edge.source.resource_name, |
| 728 | + edge.source.column_name, |
| 729 | + edge.target.resource_name, |
| 730 | + edge.target.column_name, |
| 731 | + ) |
| 732 | + for edge in result.edges |
| 733 | + ) |
| 734 | + ) |
| 735 | + assert len(set(edges_by_project_size)) == 1 |
| 736 | + assert "orders_2.amount->orders_3.amount" in edges_by_project_size[0] |
0 commit comments