diff --git a/integration_tests/tests/test_dbt_artifacts/test_project_filtering.py b/integration_tests/tests/test_dbt_artifacts/test_project_filtering.py new file mode 100644 index 000000000..0908c2da2 --- /dev/null +++ b/integration_tests/tests/test_dbt_artifacts/test_project_filtering.py @@ -0,0 +1,107 @@ +""" +Integration tests for the upload_only_current_project_artifacts config var. + +When enabled, artifact uploads should only include resources from the current +project (by package_name), excluding artifacts from dependency packages. +When disabled (default), all artifacts including dependencies should be uploaded. +""" + +import uuid + +from dbt_project import DbtProject + +TEST_MODEL = "one" + + +def test_default_includes_dependency_artifacts(dbt_project: DbtProject): + """ + By default (upload_only_current_project_artifacts=false), artifacts from + dependency packages (like 'elementary') should be present in dbt_models. + """ + dbt_project.dbt_runner.vars["disable_dbt_artifacts_autoupload"] = False + dbt_project.dbt_runner.vars["cache_artifacts"] = False + + dbt_project.dbt_runner.run(select=TEST_MODEL) + + all_models = dbt_project.read_table("dbt_models", raise_if_empty=True) + package_names = {row["package_name"] for row in all_models} + + assert "elementary" in package_names, ( + "Expected 'elementary' package artifacts to be present by default, " + f"but only found packages: {package_names}" + ) + assert "elementary_tests" in package_names, ( + "Expected 'elementary_tests' package artifacts to be present, " + f"but only found packages: {package_names}" + ) + + +def test_filtering_excludes_dependency_artifacts(dbt_project: DbtProject): + """ + When upload_only_current_project_artifacts=true, only artifacts from the + current project should be uploaded — dependency packages like 'elementary' + should be excluded. + + We first clear the dbt_models table, then run with filtering enabled via + the on-run-end hook. The diff upload sees an empty table and inserts all + (filtered) artifacts, so only current-project rows end up in the table. + """ + # Clear existing rows so the diff upload will insert fresh filtered data. + dbt_project.run_query("DELETE FROM {{ ref('dbt_models') }} WHERE 1=1") + + dbt_project.dbt_runner.vars["disable_dbt_artifacts_autoupload"] = False + dbt_project.dbt_runner.vars["upload_only_current_project_artifacts"] = True + + dbt_project.dbt_runner.run(select=TEST_MODEL) + + all_models = dbt_project.read_table("dbt_models", raise_if_empty=True) + package_names = {row["package_name"] for row in all_models} + + assert package_names == {"elementary_tests"}, ( + "Expected only 'elementary_tests' artifacts when filtering is enabled, " + f"but found packages: {package_names}" + ) + + +def test_filtering_applies_to_tests(dbt_project: DbtProject): + """ + When upload_only_current_project_artifacts=true, tests from the current + project should still be uploaded. + """ + unique_id = str(uuid.uuid4()).replace("-", "_") + model_name = f"filter_test_model_{unique_id}" + model_sql = "select 1 as col" + schema_yaml = { + "version": 2, + "models": [ + { + "name": model_name, + "columns": [{"name": "col", "tests": ["unique"]}], + } + ], + } + + with dbt_project.write_yaml( + schema_yaml, name=f"schema_filter_test_{unique_id}.yml" + ): + dbt_model_path = dbt_project.models_dir_path / "tmp" / f"{model_name}.sql" + dbt_model_path.parent.mkdir(parents=True, exist_ok=True) + dbt_model_path.write_text(model_sql) + try: + # Clear existing rows so the diff upload inserts fresh filtered data. + dbt_project.run_query("DELETE FROM {{ ref('dbt_tests') }} WHERE 1=1") + + dbt_project.dbt_runner.vars["disable_dbt_artifacts_autoupload"] = False + dbt_project.dbt_runner.vars["upload_only_current_project_artifacts"] = True + + dbt_project.dbt_runner.run(select=model_name) + + tests = dbt_project.read_table("dbt_tests", raise_if_empty=True) + test_packages = {row["package_name"] for row in tests} + assert test_packages == {"elementary_tests"}, ( + "Expected only 'elementary_tests' tests when filtering is enabled, " + f"but found packages: {test_packages}" + ) + finally: + if dbt_model_path.exists(): + dbt_model_path.unlink() diff --git a/macros/edr/dbt_artifacts/filter_to_current_project.sql b/macros/edr/dbt_artifacts/filter_to_current_project.sql new file mode 100644 index 000000000..f5f34ba00 --- /dev/null +++ b/macros/edr/dbt_artifacts/filter_to_current_project.sql @@ -0,0 +1,9 @@ +{%- macro filter_to_current_project_if_needed(entities) -%} + {%- if elementary.get_config_var("upload_only_current_project_artifacts") -%} + {% set project_name = elementary.get_project_name() %} + {% do return( + entities | selectattr("package_name", "==", project_name) | list + ) %} + {%- else -%} {% do return(entities | list) %} + {%- endif -%} +{%- endmacro -%} diff --git a/macros/edr/dbt_artifacts/upload_dbt_columns.sql b/macros/edr/dbt_artifacts/upload_dbt_columns.sql index 318d5895e..de6ca6df5 100644 --- a/macros/edr/dbt_artifacts/upload_dbt_columns.sql +++ b/macros/edr/dbt_artifacts/upload_dbt_columns.sql @@ -1,7 +1,11 @@ {%- macro upload_dbt_columns(should_commit=false, metadata_hashes=none) -%} {% set relation = elementary.get_elementary_relation("dbt_columns") %} {% if execute and relation %} - {% set tables = graph.nodes.values() | list + graph.sources.values() | list %} + {% set tables = elementary.filter_to_current_project_if_needed( + graph.nodes.values() + ) + elementary.filter_to_current_project_if_needed( + graph.sources.values() + ) %} {% do elementary.upload_artifacts_to_table( relation, tables, diff --git a/macros/edr/dbt_artifacts/upload_dbt_exposures.sql b/macros/edr/dbt_artifacts/upload_dbt_exposures.sql index 7d8949495..9200e2197 100644 --- a/macros/edr/dbt_artifacts/upload_dbt_exposures.sql +++ b/macros/edr/dbt_artifacts/upload_dbt_exposures.sql @@ -1,8 +1,9 @@ {%- macro upload_dbt_exposures(should_commit=false, metadata_hashes=none) -%} {% set relation = elementary.get_elementary_relation("dbt_exposures") %} {% if execute and relation %} - {% set exposures = graph.exposures.values() | selectattr( - "resource_type", "==", "exposure" + {% set exposures = elementary.filter_to_current_project_if_needed( + graph.exposures.values() + | selectattr("resource_type", "==", "exposure") ) %} {% do elementary.upload_artifacts_to_table( relation, diff --git a/macros/edr/dbt_artifacts/upload_dbt_models.sql b/macros/edr/dbt_artifacts/upload_dbt_models.sql index 5e68e4af5..60d36b965 100644 --- a/macros/edr/dbt_artifacts/upload_dbt_models.sql +++ b/macros/edr/dbt_artifacts/upload_dbt_models.sql @@ -1,8 +1,8 @@ {%- macro upload_dbt_models(should_commit=false, metadata_hashes=none) -%} {% set relation = elementary.get_elementary_relation("dbt_models") %} {% if execute and relation %} - {% set models = graph.nodes.values() | selectattr( - "resource_type", "==", "model" + {% set models = elementary.filter_to_current_project_if_needed( + graph.nodes.values() | selectattr("resource_type", "==", "model") ) %} {% do elementary.upload_artifacts_to_table( relation, diff --git a/macros/edr/dbt_artifacts/upload_dbt_seeds.sql b/macros/edr/dbt_artifacts/upload_dbt_seeds.sql index 1661bc491..b70f26957 100644 --- a/macros/edr/dbt_artifacts/upload_dbt_seeds.sql +++ b/macros/edr/dbt_artifacts/upload_dbt_seeds.sql @@ -1,8 +1,8 @@ {%- macro upload_dbt_seeds(should_commit=false, metadata_hashes=none) -%} {% set relation = elementary.get_elementary_relation("dbt_seeds") %} {% if execute and relation %} - {% set seeds = graph.nodes.values() | selectattr( - "resource_type", "==", "seed" + {% set seeds = elementary.filter_to_current_project_if_needed( + graph.nodes.values() | selectattr("resource_type", "==", "seed") ) %} {% do elementary.upload_artifacts_to_table( relation, diff --git a/macros/edr/dbt_artifacts/upload_dbt_snapshots.sql b/macros/edr/dbt_artifacts/upload_dbt_snapshots.sql index 6c4d414c5..963f0d39a 100644 --- a/macros/edr/dbt_artifacts/upload_dbt_snapshots.sql +++ b/macros/edr/dbt_artifacts/upload_dbt_snapshots.sql @@ -1,8 +1,9 @@ {%- macro upload_dbt_snapshots(should_commit=false, metadata_hashes=none) -%} {% set relation = elementary.get_elementary_relation("dbt_snapshots") %} {% if execute and relation %} - {% set snapshots = graph.nodes.values() | selectattr( - "resource_type", "==", "snapshot" + {% set snapshots = elementary.filter_to_current_project_if_needed( + graph.nodes.values() + | selectattr("resource_type", "==", "snapshot") ) %} {% do elementary.upload_artifacts_to_table( relation, diff --git a/macros/edr/dbt_artifacts/upload_dbt_sources.sql b/macros/edr/dbt_artifacts/upload_dbt_sources.sql index 2f770b2c7..198d67146 100644 --- a/macros/edr/dbt_artifacts/upload_dbt_sources.sql +++ b/macros/edr/dbt_artifacts/upload_dbt_sources.sql @@ -1,8 +1,9 @@ {%- macro upload_dbt_sources(should_commit=false, metadata_hashes=none) -%} {% set relation = elementary.get_elementary_relation("dbt_sources") %} {% if execute and relation %} - {% set sources = graph.sources.values() | selectattr( - "resource_type", "==", "source" + {% set sources = elementary.filter_to_current_project_if_needed( + graph.sources.values() + | selectattr("resource_type", "==", "source") ) %} {% do elementary.upload_artifacts_to_table( relation, diff --git a/macros/edr/dbt_artifacts/upload_dbt_tests.sql b/macros/edr/dbt_artifacts/upload_dbt_tests.sql index aabd17156..78593581d 100644 --- a/macros/edr/dbt_artifacts/upload_dbt_tests.sql +++ b/macros/edr/dbt_artifacts/upload_dbt_tests.sql @@ -1,8 +1,8 @@ {%- macro upload_dbt_tests(should_commit=false, metadata_hashes=none) -%} {% set relation = elementary.get_elementary_relation("dbt_tests") %} {% if execute and relation %} - {% set tests = graph.nodes.values() | selectattr( - "resource_type", "==", "test" + {% set tests = elementary.filter_to_current_project_if_needed( + graph.nodes.values() | selectattr("resource_type", "==", "test") ) %} {% do elementary.upload_artifacts_to_table( relation, diff --git a/macros/edr/system/system_utils/get_config_var.sql b/macros/edr/system/system_utils/get_config_var.sql index 261c431b7..186901a7e 100644 --- a/macros/edr/system/system_utils/get_config_var.sql +++ b/macros/edr/system/system_utils/get_config_var.sql @@ -144,6 +144,7 @@ "disable_samples_on_pii_tags": false, "pii_tags": ["pii"], "bigquery_disable_partitioning": false, + "upload_only_current_project_artifacts": false, } %} {{- return(default_config) -}} {%- endmacro -%}