Skip to content
Merged
Show file tree
Hide file tree
Changes from 5 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 106 additions & 0 deletions integration_tests/tests/test_dbt_artifacts/test_project_filtering.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
"""
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)
Comment thread
coderabbitai[bot] marked this conversation as resolved.
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 run the dbt_models model directly (which does a full table replace via
post_hook) to ensure a clean slate, rather than relying on the on-run-end
diff upload which may not remove pre-existing rows.
"""
dbt_project.dbt_runner.vars["disable_dbt_artifacts_autoupload"] = True
dbt_project.dbt_runner.vars["cache_artifacts"] = False
dbt_project.dbt_runner.vars["upload_only_current_project_artifacts"] = True

# Running dbt_models directly triggers its post_hook (upload_dbt_models),
# which does a full table replace with the filtered artifacts.
dbt_project.dbt_runner.run(select="dbt_models")

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, tmp_path):
Comment thread
coderabbitai[bot] marked this conversation as resolved.
Outdated
"""
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:
dbt_project.dbt_runner.vars["disable_dbt_artifacts_autoupload"] = True
dbt_project.dbt_runner.vars["cache_artifacts"] = False
dbt_project.dbt_runner.vars["upload_only_current_project_artifacts"] = True

# Run dbt_tests model directly to trigger post_hook with full replace
dbt_project.dbt_runner.run(select=f"{model_name} dbt_tests")

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()
9 changes: 9 additions & 0 deletions macros/edr/dbt_artifacts/filter_to_current_project.sql
Original file line number Diff line number Diff line change
@@ -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 -%}
6 changes: 5 additions & 1 deletion macros/edr/dbt_artifacts/upload_dbt_columns.sql
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
5 changes: 3 additions & 2 deletions macros/edr/dbt_artifacts/upload_dbt_exposures.sql
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
4 changes: 2 additions & 2 deletions macros/edr/dbt_artifacts/upload_dbt_models.sql
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
4 changes: 2 additions & 2 deletions macros/edr/dbt_artifacts/upload_dbt_seeds.sql
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
5 changes: 3 additions & 2 deletions macros/edr/dbt_artifacts/upload_dbt_snapshots.sql
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
5 changes: 3 additions & 2 deletions macros/edr/dbt_artifacts/upload_dbt_sources.sql
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
4 changes: 2 additions & 2 deletions macros/edr/dbt_artifacts/upload_dbt_tests.sql
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
1 change: 1 addition & 0 deletions macros/edr/system/system_utils/get_config_var.sql
Original file line number Diff line number Diff line change
Expand Up @@ -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 -%}
Expand Down
Loading