Refactor LoadMethod.LOCAL to use symlinks instead of copying directory

cosmos/dbt/graph.py

@@ -25,6 +25,7 @@
 from cosmos.dbt.parser.project import LegacyDbtProject
 from cosmos.dbt.selector import select_nodes
 from cosmos.log import get_logger
+from cosmos.utils import create_symlinks
 
 logger = get_logger(__name__)
 
@@ -53,14 +54,6 @@ class DbtNode:
     has_test: bool = False
 
 
-def create_symlinks(project_path: Path, tmp_dir: Path) -> None:
-    """Helper function to create symlinks to the dbt project files."""
-    ignore_paths = (DBT_LOG_DIR_NAME, DBT_TARGET_DIR_NAME, "dbt_packages", "profiles.yml")
-    for child_name in os.listdir(project_path):
-        if child_name not in ignore_paths:
-            os.symlink(project_path / child_name, tmp_dir / child_name)
-
-
 def run_command(command: list[str], tmp_dir: Path, env_vars: dict[str, str]) -> str:
     """Run a command in a subprocess, returning the stdout."""
     logger.info("Running command: `%s`", " ".join(command))

cosmos/operators/local.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import os
-import shutil
 import signal
 import tempfile
 from attr import define
@@ -45,6 +44,7 @@
     FullOutputSubprocessResult,
 )
 from cosmos.dbt.parser.output import extract_log_issues, parse_output
+from cosmos.utils import create_symlinks
 
 DBT_NO_TESTS_MSG = "Nothing to do"
 DBT_WARN_MSG = "WARN"
@@ -192,19 +192,14 @@ def run_command(
         """
         Copies the dbt project to a temporary directory and runs the command.
         """
-        with tempfile.TemporaryDirectory() as tmp_dir:
+        with tempfile.TemporaryDirectory() as tmp_project_dir:
             logger.info(
                 "Cloning project to writable temp directory %s from %s",
-                tmp_dir,
+                tmp_project_dir,
                 self.project_dir,
             )
 
-            # need a subfolder because shutil.copytree will fail if the destination dir already exists
-            tmp_project_dir = os.path.join(tmp_dir, "dbt_project")
-            shutil.copytree(
-                self.project_dir,
-                tmp_project_dir,
-            )
+            create_symlinks(Path(self.project_dir), Path(tmp_project_dir))
 
             # if we need to install deps, do so
             if self.install_deps:
@@ -397,6 +392,8 @@ class DbtSeedLocalOperator(DbtLocalBaseOperator):
 
     ui_color = "#F58D7E"
 
+    template_fields: Sequence[str] = DbtBaseOperator.template_fields + ("full_refresh",)  # type: ignore[operator]
+
     def __init__(self, full_refresh: bool = False, **kwargs: Any) -> None:
         self.full_refresh = full_refresh
         super().__init__(**kwargs)
@@ -434,6 +431,7 @@ class DbtRunLocalOperator(DbtLocalBaseOperator):
 
     ui_color = "#7352BA"
     ui_fgcolor = "#F4F2FC"
+    template_fields: Sequence[str] = DbtBaseOperator.template_fields + ("full_refresh",)  # type: ignore[operator]
 
     def __init__(self, full_refresh: bool = False, **kwargs: Any) -> None:
         self.full_refresh = full_refresh

cosmos/profiles/__init__.py

@@ -16,7 +16,8 @@
 from .redshift.user_pass import RedshiftUserPasswordProfileMapping
 from .snowflake.user_pass import SnowflakeUserPasswordProfileMapping
 from .snowflake.user_privatekey import SnowflakePrivateKeyPemProfileMapping
-from .snowflake.user_encrypted_privatekey import SnowflakeEncryptedPrivateKeyPemProfileMapping
+from .snowflake.user_encrypted_privatekey_file import SnowflakeEncryptedPrivateKeyFilePemProfileMapping
+from .snowflake.user_encrypted_privatekey_env_variable import SnowflakeEncryptedPrivateKeyPemProfileMapping
 from .spark.thrift import SparkThriftProfileMapping
 from .trino.certificate import TrinoCertificateProfileMapping
 from .trino.jwt import TrinoJWTProfileMapping
@@ -32,6 +33,7 @@
     PostgresUserPasswordProfileMapping,
     RedshiftUserPasswordProfileMapping,
     SnowflakeUserPasswordProfileMapping,
+    SnowflakeEncryptedPrivateKeyFilePemProfileMapping,
     SnowflakeEncryptedPrivateKeyPemProfileMapping,
     SnowflakePrivateKeyPemProfileMapping,
     SparkThriftProfileMapping,
@@ -71,6 +73,7 @@ def get_automatic_profile_mapping(
     "RedshiftUserPasswordProfileMapping",
     "SnowflakeUserPasswordProfileMapping",
     "SnowflakePrivateKeyPemProfileMapping",
+    "SnowflakeEncryptedPrivateKeyFilePemProfileMapping",
     "SparkThriftProfileMapping",
     "ExasolUserPasswordProfileMapping",
     "TrinoLDAPProfileMapping",

cosmos/profiles/snowflake/__init__.py

@@ -2,10 +2,12 @@
 
 from .user_pass import SnowflakeUserPasswordProfileMapping
 from .user_privatekey import SnowflakePrivateKeyPemProfileMapping
-from .user_encrypted_privatekey import SnowflakeEncryptedPrivateKeyPemProfileMapping
+from .user_encrypted_privatekey_file import SnowflakeEncryptedPrivateKeyFilePemProfileMapping
+from .user_encrypted_privatekey_env_variable import SnowflakeEncryptedPrivateKeyPemProfileMapping
 
 __all__ = [
     "SnowflakeUserPasswordProfileMapping",
     "SnowflakePrivateKeyPemProfileMapping",
+    "SnowflakeEncryptedPrivateKeyFilePemProfileMapping",
     "SnowflakeEncryptedPrivateKeyPemProfileMapping",
 ]

cosmos/profiles/snowflake/user_encrypted_privatekey_env_variable.py

@@ -0,0 +1,93 @@
+"Maps Airflow Snowflake connections to dbt profiles if they use a user/private key."
+from __future__ import annotations
+
+import json
+from typing import TYPE_CHECKING, Any
+
+from ..base import BaseProfileMapping
+
+if TYPE_CHECKING:
+    from airflow.models import Connection
+
+
+class SnowflakeEncryptedPrivateKeyPemProfileMapping(BaseProfileMapping):
+    """
+    Maps Airflow Snowflake connections to dbt profiles if they use a user/private key.
+    https://docs.getdbt.com/docs/core/connect-data-platform/snowflake-setup#key-pair-authentication
+    https://airflow.apache.org/docs/apache-airflow-providers-snowflake/stable/connections/snowflake.html
+    """
+
+    airflow_connection_type: str = "snowflake"
+    dbt_profile_type: str = "snowflake"
+    is_community: bool = True
+
+    required_fields = [
+        "account",
+        "user",
+        "database",
+        "warehouse",
+        "schema",
+        "private_key",
+        "private_key_passphrase",
+    ]
+    secret_fields = [
+        "private_key",
+        "private_key_passphrase",
+    ]
+    airflow_param_mapping = {
+        "account": "extra.account",
+        "user": "login",
+        "database": "extra.database",
+        "warehouse": "extra.warehouse",
+        "schema": "schema",
+        "role": "extra.role",
+        "private_key": "extra.private_key_content",
+        "private_key_passphrase": "password",
+    }
+
+    def can_claim_connection(self) -> bool:
+        # Make sure this isn't a private key path credential
+        result = super().can_claim_connection()
+        if result and self.conn.extra_dejson.get("private_key_file") is not None:
+            return False
+        return result
+
+    @property
+    def conn(self) -> Connection:
+        """
+        Snowflake can be odd because the fields used to be stored with keys in the format
+        'extra__snowflake__account', but now are stored as 'account'.
+
+        This standardizes the keys to be 'account', 'database', etc.
+        """
+        conn = super().conn
+
+        conn_dejson = conn.extra_dejson
+
+        if conn_dejson.get("extra__snowflake__account"):
+            conn_dejson = {key.replace("extra__snowflake__", ""): value for key, value in conn_dejson.items()}
+
+        conn.extra = json.dumps(conn_dejson)
+
+        return conn
+
+    @property
+    def profile(self) -> dict[str, Any | None]:
+        "Gets profile."
+        profile_vars = {
+            **self.mapped_params,
+            **self.profile_args,
+            "private_key": self.get_env_var_format("private_key"),
+            "private_key_passphrase": self.get_env_var_format("private_key_passphrase"),
+        }
+
+        # remove any null values
+        return self.filter_null(profile_vars)
+
+    def transform_account(self, account: str) -> str:
+        "Transform the account to the format <account>.<region> if it's not already."
+        region = self.conn.extra_dejson.get("region")
+        if region and region not in account:
+            account = f"{account}.{region}"
+
+        return str(account)

cosmos/profiles/snowflake/user_encrypted_privatekey.py → cosmos/profiles/snowflake/user_encrypted_privatekey_file.py

@@ -1,4 +1,4 @@
-"Maps Airflow Snowflake connections to dbt profiles if they use a user/private key."
+"Maps Airflow Snowflake connections to dbt profiles if they use a user/private key path."
 from __future__ import annotations
 
 import json
@@ -10,9 +10,9 @@
     from airflow.models import Connection
 
 
-class SnowflakeEncryptedPrivateKeyPemProfileMapping(BaseProfileMapping):
+class SnowflakeEncryptedPrivateKeyFilePemProfileMapping(BaseProfileMapping):
     """
-    Maps Airflow Snowflake connections to dbt profiles if they use a user/private key.
+    Maps Airflow Snowflake connections to dbt profiles if they use a user/private key path.
     https://docs.getdbt.com/docs/core/connect-data-platform/snowflake-setup#key-pair-authentication
     https://airflow.apache.org/docs/apache-airflow-providers-snowflake/stable/connections/snowflake.html
     """
@@ -44,6 +44,13 @@ class SnowflakeEncryptedPrivateKeyPemProfileMapping(BaseProfileMapping):
         "private_key_path": "extra.private_key_file",
     }
 
+    def can_claim_connection(self) -> bool:
+        # Make sure this isn't a private key environmentvariable
+        result = super().can_claim_connection()
+        if result and self.conn.extra_dejson.get("private_key_content") is not None:
+            return False
+        return result
+
     @property
     def conn(self) -> Connection:
         """

cosmos/profiles/snowflake/user_pass.py

@@ -44,7 +44,10 @@ class SnowflakeUserPasswordProfileMapping(BaseProfileMapping):
     def can_claim_connection(self) -> bool:
         # Make sure this isn't a private key path credential
         result = super().can_claim_connection()
-        if result and self.conn.extra_dejson.get("private_key_file") is not None:
+        if result and (
+            self.conn.extra_dejson.get("private_key_file") is not None
+            or self.conn.extra_dejson.get("private_key_content") is not None
+        ):
             return False
         return result
 

cosmos/utils.py

@@ -0,0 +1,14 @@
+from pathlib import Path
+import os
+from cosmos.constants import (
+    DBT_LOG_DIR_NAME,
+    DBT_TARGET_DIR_NAME,
+)
+
+
+def create_symlinks(project_path: Path, tmp_dir: Path) -> None:
+    """Helper function to create symlinks to the dbt project files."""
+    ignore_paths = (DBT_LOG_DIR_NAME, DBT_TARGET_DIR_NAME, "dbt_packages", "profiles.yml")
+    for child_name in os.listdir(project_path):
+        if child_name not in ignore_paths:
+            os.symlink(project_path / child_name, tmp_dir / child_name)

tests/dbt/test_graph.py

@@ -12,7 +12,6 @@
     DbtGraph,
     DbtNode,
     LoadMode,
-    create_symlinks,
     run_command,
     parse_dbt_ls_output,
 )
@@ -659,17 +658,6 @@ def test_load_dbt_ls_and_manifest_with_model_version(load_method):
     } == set(dbt_graph.nodes["model.jaffle_shop.orders"].depends_on)
 
 
-def test_create_symlinks(tmp_path):
-    """Tests that symlinks are created for expected files in the dbt project directory."""
-    tmp_dir = tmp_path / "dbt-project"
-    tmp_dir.mkdir()
-
-    create_symlinks(DBT_PROJECTS_ROOT_DIR / "jaffle_shop", tmp_dir)
-    for child in tmp_dir.iterdir():
-        assert child.is_symlink()
-        assert child.name not in ("logs", "target", "profiles.yml", "dbt_packages")
-
-
 @pytest.mark.parametrize(
     "stdout,returncode",
     [

tests/operators/test_local.py

@@ -420,6 +420,19 @@ def test_calculate_openlineage_events_completes_openlineage_errors(mock_processo
     assert err_msg in caplog.text
 
 
+@pytest.mark.parametrize(
+    "operator_class,expected_template",
+    [
+        (DbtSeedLocalOperator, ("env", "vars", "full_refresh")),
+        (DbtRunLocalOperator, ("env", "vars", "full_refresh")),
+    ],
+)
+def test_dbt_base_operator_template_fields(operator_class, expected_template):
+    # Check if value of template fields is what we expect for the operators we're validating
+    dbt_base_operator = operator_class(profile_config=profile_config, task_id="my-task", project_dir="my/dir")
+    assert dbt_base_operator.template_fields == expected_template
+
+
 @patch.object(DbtDocsGCSLocalOperator, "required_files", ["file1", "file2"])
 def test_dbt_docs_gcs_local_operator():
     mock_gcs = MagicMock()

tests/profiles/snowflake/test_snowflake_user_encrypted_privatekey.py → tests/profiles/snowflake/test_snowflake_user_encrypted_privatekey_env_variable.py

@@ -1,4 +1,4 @@
-"Tests for the Snowflake user/private key profile."
+"Tests for the Snowflake user/private key environmentvariable profile."
 
 import json
 from unittest.mock import patch
@@ -29,7 +29,7 @@ def mock_snowflake_conn():  # type: ignore
                 "region": "my_region",
                 "database": "my_database",
                 "warehouse": "my_warehouse",
-                "private_key_file": "path/to/private_key.p8",
+                "private_key_content": "my_private_key",
             }
         ),
     )
@@ -52,7 +52,7 @@ def test_connection_claiming() -> None:
                 "account": "my_account",
                 "database": "my_database",
                 "warehouse": "my_warehouse",
-                "private_key_file": "path/to/private_key.p8",
+                "private_key_content": "my_private_key",
             }
         ),
     }
@@ -130,8 +130,8 @@ def test_profile_args(
     assert profile_mapping.profile == {
         "type": mock_snowflake_conn.conn_type,
         "user": mock_snowflake_conn.login,
+        "private_key": "{{ env_var('COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY') }}",
         "private_key_passphrase": "{{ env_var('COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY_PASSPHRASE') }}",
-        "private_key_path": mock_snowflake_conn.extra_dejson.get("private_key_file"),
         "schema": mock_snowflake_conn.schema,
         "account": f"{mock_account}.{mock_region}",
         "database": mock_snowflake_conn.extra_dejson.get("database"),
@@ -160,7 +160,7 @@ def test_profile_args_overrides(
         "type": mock_snowflake_conn.conn_type,
         "user": mock_snowflake_conn.login,
         "private_key_passphrase": "{{ env_var('COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY_PASSPHRASE') }}",
-        "private_key_path": mock_snowflake_conn.extra_dejson.get("private_key_file"),
+        "private_key": "{{ env_var('COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY') }}",
         "schema": mock_snowflake_conn.schema,
         "account": f"{mock_account}.{mock_region}",
         "database": "my_db_override",
@@ -178,6 +178,7 @@ def test_profile_env_vars(
         mock_snowflake_conn.conn_id,
     )
     assert profile_mapping.env_vars == {
+        "COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY": mock_snowflake_conn.extra_dejson.get("private_key_content"),
         "COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY_PASSPHRASE": mock_snowflake_conn.password,
     }
 
@@ -197,7 +198,7 @@ def test_old_snowflake_format() -> None:
                 "extra__snowflake__account": "my_account",
                 "extra__snowflake__database": "my_database",
                 "extra__snowflake__warehouse": "my_warehouse",
-                "extra__snowflake__private_key_file": "path/to/private_key.p8",
+                "extra__snowflake__private_key_content": "my_private_key",
             }
         ),
     )
@@ -207,8 +208,8 @@ def test_old_snowflake_format() -> None:
         assert profile_mapping.profile == {
             "type": conn.conn_type,
             "user": conn.login,
+            "private_key": "{{ env_var('COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY') }}",
             "private_key_passphrase": "{{ env_var('COSMOS_CONN_SNOWFLAKE_PRIVATE_KEY_PASSPHRASE') }}",
-            "private_key_path": conn.extra_dejson.get("private_key_file"),
             "schema": conn.schema,
             "account": conn.extra_dejson.get("account"),
             "database": conn.extra_dejson.get("database"),