Skip to content

Commit

Permalink
Fix LoadMode.AUTOMATIC behaviour to use LoadMode.DBT_LS when `Pro…
Browse files Browse the repository at this point in the history
…fileMapping` is used (#625)

Since #489 was merged, the behavior of `LoadMode.AUTOMATIC` changed to
generate a `profiles.yml` file if the file didn't exist. However, we
forgot to remove the previously necessary condition for being able to
run `LoadMode.DBT_LS` (having the `profiles.yml` file).

This leads to inconsistent behaviour in Cosmos when using
`LoadMode.AUTOMATIC` and the `manifest.json` was not available:
1. If the user used a `ProfileConfig` with `profiles_yml_filepath`, it
would use `LoadMode.DBT_LS`
2. If the user used a `ProfileConfig` with a ProfileMapping class, it
would unnecessarily use `LoadMode.CUSTOM`

This PR fixes the behaviour to attempt to use `LoadMode.DBT_LS`
regardless of how the `ProfileConfig` was set.
  • Loading branch information
tatiana authored Oct 25, 2023
1 parent a4aa5cc commit ad7dcf0
Show file tree
Hide file tree
Showing 4 changed files with 24 additions and 17 deletions.
6 changes: 0 additions & 6 deletions cosmos/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -163,12 +163,6 @@ def validate_profile(self) -> None:
if not self.profiles_yml_filepath and not self.profile_mapping:
raise CosmosValueError("Either profiles_yml_filepath or profile_mapping must be set to render a profile")

def is_profile_yml_available(self) -> bool:
"""
Check if the `dbt` profiles.yml file exists.
"""
return Path(self.profiles_yml_filepath).exists() if self.profiles_yml_filepath else False

@contextlib.contextmanager
def ensure_profile(
self, desired_profile_path: Path | None = None, use_mock_values: bool = False
Expand Down
6 changes: 1 addition & 5 deletions cosmos/dbt/graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -121,11 +121,7 @@ def load(
if self.project.is_manifest_available():
self.load_from_dbt_manifest()
else:
if (
execution_mode == ExecutionMode.LOCAL
and self.profile_config
and self.profile_config.is_profile_yml_available()
):
if execution_mode == ExecutionMode.LOCAL and self.profile_config:
try:
self.load_via_dbt_ls()
except FileNotFoundError:
Expand Down
26 changes: 21 additions & 5 deletions tests/dbt/test_graph.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,12 +82,10 @@ def test_load_automatic_manifest_is_available(mock_load_from_dbt_manifest):
assert mock_load_from_dbt_manifest.called


@patch("cosmos.dbt.graph.DbtGraph.load_via_custom_parser", side_effect=FileNotFoundError())
@patch("cosmos.dbt.graph.DbtGraph.load_via_custom_parser", side_effect=None)
@patch("cosmos.dbt.graph.DbtGraph.load_via_dbt_ls", return_value=None)
def test_load_automatic_without_manifest(mock_load_via_dbt_ls, mock_load_via_custom_parser):
project_config = ProjectConfig(
dbt_project_path=DBT_PROJECTS_ROOT_DIR / DBT_PROJECT_NAME, manifest_path="/tmp/manifest.json"
)
def test_load_automatic_without_manifest_with_profile_yml(mock_load_via_dbt_ls, mock_load_via_custom_parser):
project_config = ProjectConfig(dbt_project_path=DBT_PROJECTS_ROOT_DIR / DBT_PROJECT_NAME)
profile_config = ProfileConfig(
profile_name="test",
target_name="test",
Expand All @@ -99,6 +97,24 @@ def test_load_automatic_without_manifest(mock_load_via_dbt_ls, mock_load_via_cus
assert not mock_load_via_custom_parser.called


@patch("cosmos.dbt.graph.DbtGraph.load_via_custom_parser", side_effect=None)
@patch("cosmos.dbt.graph.DbtGraph.load_via_dbt_ls", return_value=None)
def test_load_automatic_without_manifest_with_profile_mapping(mock_load_via_dbt_ls, mock_load_via_custom_parser):
project_config = ProjectConfig(dbt_project_path=DBT_PROJECTS_ROOT_DIR / DBT_PROJECT_NAME)
profile_config = ProfileConfig(
profile_name="test",
target_name="test",
profile_mapping=PostgresUserPasswordProfileMapping(
conn_id="airflow_db",
profile_args={"schema": "public"},
),
)
dbt_graph = DbtGraph(project=project_config, profile_config=profile_config)
dbt_graph.load(execution_mode=ExecutionMode.LOCAL)
assert mock_load_via_dbt_ls.called
assert not mock_load_via_custom_parser.called


@patch("cosmos.dbt.graph.DbtGraph.load_via_custom_parser", return_value=None)
@patch("cosmos.dbt.graph.DbtGraph.load_via_dbt_ls", side_effect=FileNotFoundError())
def test_load_automatic_without_manifest_and_without_dbt_cmd(mock_load_via_dbt_ls, mock_load_via_custom_parser):
Expand Down
3 changes: 2 additions & 1 deletion tests/operators/test_local.py
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,8 @@ def test_dbt_base_operator_use_indirect_selection(indirect_selection_type) -> No
assert cmd[-2] == "--indirect-selection"
assert cmd[-1] == indirect_selection_type
else:
assert cmd == ["dbt", "run"]
assert cmd[0].endswith("dbt")
assert cmd[1] == "run"


@pytest.mark.parametrize(
Expand Down

0 comments on commit ad7dcf0

Please sign in to comment.