From ad7dcf0b419f247cab7effcc205ce16a9b922d32 Mon Sep 17 00:00:00 2001
From: Tatiana Al-Chueyr <tatiana.alchueyr@gmail.com>
Date: Wed, 25 Oct 2023 19:10:10 +0100
Subject: [PATCH] Fix `LoadMode.AUTOMATIC` behaviour to use `LoadMode.DBT_LS`
 when `ProfileMapping` 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.
---
 cosmos/config.py              |  6 ------
 cosmos/dbt/graph.py           |  6 +-----
 tests/dbt/test_graph.py       | 26 +++++++++++++++++++++-----
 tests/operators/test_local.py |  3 ++-
 4 files changed, 24 insertions(+), 17 deletions(-)

diff --git a/cosmos/config.py b/cosmos/config.py
index 15f88e037..610e6e489 100644
--- a/cosmos/config.py
+++ b/cosmos/config.py
@@ -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
diff --git a/cosmos/dbt/graph.py b/cosmos/dbt/graph.py
index ac92c46f5..c41ee0713 100644
--- a/cosmos/dbt/graph.py
+++ b/cosmos/dbt/graph.py
@@ -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:
diff --git a/tests/dbt/test_graph.py b/tests/dbt/test_graph.py
index f176fa445..3927bbfdd 100644
--- a/tests/dbt/test_graph.py
+++ b/tests/dbt/test_graph.py
@@ -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",
@@ -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):
diff --git a/tests/operators/test_local.py b/tests/operators/test_local.py
index b883adea5..580d49e6c 100644
--- a/tests/operators/test_local.py
+++ b/tests/operators/test_local.py
@@ -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(