Jw/snow 2150049 dbt external access integrations (#2474)

* feat: [SNOW-2150049] dbt deploy - add external access integration support

* refactor: [SNOW-2150049] rename

Author: sfc-gh-jwilkowski

src/snowflake/cli/_plugins/dbt/commands.py

@@ -111,6 +111,13 @@ def deploy_dbt(
         help="Unset the default target for the dbt project. Mutually exclusive with --default-target.",
         hidden=FeatureFlag.ENABLE_DBT_GA_FEATURES.is_disabled(),
     ),
+    external_access_integrations: Optional[list[str]] = typer.Option(
+        None,
+        "--external-access-integration",
+        show_default=False,
+        help="External access integration to be used by the dbt object.",
+        hidden=FeatureFlag.ENABLE_DBT_GA_FEATURES.is_disabled(),
+    ),
     **options,
 ) -> CommandResult:
     """
@@ -121,6 +128,7 @@ def deploy_dbt(
     if FeatureFlag.ENABLE_DBT_GA_FEATURES.is_disabled():
         default_target = None
         unset_default_target = False
+        external_access_integrations = None
 
     project_path = SecurePath(source) if source is not None else SecurePath.cwd()
     profiles_dir_path = SecurePath(profiles_dir) if profiles_dir else project_path
@@ -132,6 +140,7 @@ def deploy_dbt(
             force=force,
             default_target=default_target,
             unset_default_target=unset_default_target,
+            external_access_integrations=external_access_integrations,
         )
     )
 

src/snowflake/cli/_plugins/dbt/manager.py

@@ -17,7 +17,7 @@
 from collections import defaultdict
 from pathlib import Path
 from tempfile import TemporaryDirectory
-from typing import Optional, TypedDict
+from typing import List, Optional, TypedDict
 
 import yaml
 from snowflake.cli._plugins.dbt.constants import PROFILES_FILENAME
@@ -85,6 +85,7 @@ def deploy(
         force: bool,
         default_target: Optional[str] = None,
         unset_default_target: bool = False,
+        external_access_integrations: Optional[List[str]] = None,
     ) -> SnowflakeCursor:
         dbt_project_path = path / "dbt_project.yml"
         if not dbt_project_path.exists():
@@ -123,38 +124,95 @@ def deploy(
 
         with cli_console.phase("Creating DBT project"):
             if force is True:
-                query = f"CREATE OR REPLACE DBT PROJECT {fqn}"
-                query += f"\nFROM {stage_name}"
-                if default_target:
-                    query += f" DEFAULT_TARGET='{default_target}'"
-                return self.execute_query(query)
+                return self._deploy_create_or_replace(
+                    fqn, stage_name, default_target, external_access_integrations
+                )
             else:
                 dbt_object_attributes = self.get_dbt_object_attributes(fqn)
                 if dbt_object_attributes is not None:
-                    # Project exists - add new version
-                    query = f"ALTER DBT PROJECT {fqn} ADD VERSION"
-                    query += f"\nFROM {stage_name}"
-                    result = self.execute_query(query)
+                    return self._deploy_alter(
+                        fqn,
+                        stage_name,
+                        dbt_object_attributes,
+                        default_target,
+                        unset_default_target,
+                        external_access_integrations,
+                    )
+                else:
+                    return self._deploy_create(
+                        fqn, stage_name, default_target, external_access_integrations
+                    )
 
-                    current_default_target = dbt_object_attributes.get("default_target")
-                    if unset_default_target and current_default_target is not None:
-                        unset_query = f"ALTER DBT PROJECT {fqn} UNSET DEFAULT_TARGET"
-                        self.execute_query(unset_query)
-                    elif default_target and (
-                        current_default_target is None
-                        or current_default_target.lower() != default_target.lower()
-                    ):
-                        set_default_query = f"ALTER DBT PROJECT {fqn} SET DEFAULT_TARGET='{default_target}'"
-                        self.execute_query(set_default_query)
+    def _deploy_alter(
+        self,
+        fqn: FQN,
+        stage_name: str,
+        dbt_object_attributes: DBTObjectEditableAttributes,
+        default_target: Optional[str],
+        unset_default_target: bool,
+        external_access_integrations: Optional[List[str]],
+    ) -> SnowflakeCursor:
+        query = f"ALTER DBT PROJECT {fqn} ADD VERSION"
+        query += f"\nFROM {stage_name}"
+        query = self._handle_external_access_integrations_query(
+            query, external_access_integrations
+        )
+        result = self.execute_query(query)
+        current_default_target = dbt_object_attributes.get("default_target")
+        if unset_default_target and current_default_target is not None:
+            unset_query = f"ALTER DBT PROJECT {fqn} UNSET DEFAULT_TARGET"
+            self.execute_query(unset_query)
+        elif default_target and (
+            current_default_target is None
+            or current_default_target.lower() != default_target.lower()
+        ):
+            set_default_query = (
+                f"ALTER DBT PROJECT {fqn} SET DEFAULT_TARGET='{default_target}'"
+            )
+            self.execute_query(set_default_query)
+        return result
 
-                    return result
-                else:
-                    # Project doesn't exist - create new one
-                    query = f"CREATE DBT PROJECT {fqn}"
-                    query += f"\nFROM {stage_name}"
-                    if default_target:
-                        query += f" DEFAULT_TARGET='{default_target}'"
-                    return self.execute_query(query)
+    def _deploy_create(
+        self,
+        fqn: FQN,
+        stage_name: str,
+        default_target: Optional[str],
+        external_access_integrations: Optional[List[str]],
+    ) -> SnowflakeCursor:
+        # Project doesn't exist - create new one
+        query = f"CREATE DBT PROJECT {fqn}"
+        query += f"\nFROM {stage_name}"
+        if default_target:
+            query += f" DEFAULT_TARGET='{default_target}'"
+        query = self._handle_external_access_integrations_query(
+            query, external_access_integrations
+        )
+        return self.execute_query(query)
+
+    @staticmethod
+    def _handle_external_access_integrations_query(
+        query: str, external_access_integrations: Optional[List[str]]
+    ) -> str:
+        if external_access_integrations:
+            integrations_str = ", ".join(external_access_integrations)
+            query += f"\nEXTERNAL_ACCESS_INTEGRATIONS = ({integrations_str})"
+        return query
+
+    def _deploy_create_or_replace(
+        self,
+        fqn: FQN,
+        stage_name: str,
+        default_target: Optional[str],
+        external_access_integrations: Optional[List[str]],
+    ) -> SnowflakeCursor:
+        query = f"CREATE OR REPLACE DBT PROJECT {fqn}"
+        query += f"\nFROM {stage_name}"
+        if default_target:
+            query += f" DEFAULT_TARGET='{default_target}'"
+        query = self._handle_external_access_integrations_query(
+            query, external_access_integrations
+        )
+        return self.execute_query(query)
 
     @staticmethod
     def _validate_profiles(

tests/dbt/test_dbt_commands.py

@@ -26,6 +26,8 @@
 )
 from snowflake.cli._plugins.dbt.manager import DBTObjectEditableAttributes
 from snowflake.cli.api.feature_flags import FeatureFlag
+from snowflake.cli.api.identifiers import FQN
+from snowflake.cli.api.secure_path import SecurePath
 
 from tests_common.feature_flag_utils import with_feature_flags
 
@@ -125,6 +127,13 @@ def mock_from_resource(self):
         ) as _fixture:
             yield _fixture
 
+    @pytest.fixture
+    def mock_deploy(self):
+        with mock.patch(
+            "snowflake.cli._plugins.dbt.manager.DBTManager.deploy"
+        ) as _fixture:
+            yield _fixture
+
     @mock.patch("snowflake.cli._plugins.dbt.manager.StageManager.put_recursive")
     @mock.patch("snowflake.cli._plugins.dbt.manager.StageManager.create")
     def test_deploys_project_from_source(
@@ -298,6 +307,68 @@ def test_deploys_project_with_fqn_uses_name_only_for_stage(
 FROM @MockDatabase.MockSchema.DBT_PROJECT_TEST_DBT_PROJECT_{mock_time()}_STAGE"""
         )
 
+    @with_feature_flags({FeatureFlag.ENABLE_DBT_GA_FEATURES: True})
+    def test_deploys_project_with_single_external_access_integration(
+        self,
+        runner,
+        dbt_project_path,
+        mock_deploy,
+    ):
+
+        result = runner.invoke(
+            [
+                "dbt",
+                "deploy",
+                "TEST_PIPELINE",
+                f"--source={dbt_project_path}",
+                "--external-access-integration",
+                "google_apis_access_integration",
+            ]
+        )
+
+        assert result.exit_code == 0, result.output
+        mock_deploy.assert_called_once_with(
+            FQN.from_string("TEST_PIPELINE"),
+            SecurePath(dbt_project_path),
+            SecurePath(dbt_project_path),
+            force=False,
+            default_target=None,
+            unset_default_target=False,
+            external_access_integrations=["google_apis_access_integration"],
+        )
+
+    @with_feature_flags({FeatureFlag.ENABLE_DBT_GA_FEATURES: True})
+    def test_deploys_project_with_multiple_external_access_integrations(
+        self,
+        runner,
+        dbt_project_path,
+        mock_deploy,
+    ):
+
+        result = runner.invoke(
+            [
+                "dbt",
+                "deploy",
+                "TEST_PIPELINE",
+                f"--source={dbt_project_path}",
+                "--external-access-integration",
+                "google_apis_access_integration",
+                "--external-access-integration",
+                "dbt_hub",
+            ]
+        )
+
+        assert result.exit_code == 0, result.output
+        mock_deploy.assert_called_once_with(
+            FQN.from_string("TEST_PIPELINE"),
+            SecurePath(dbt_project_path),
+            SecurePath(dbt_project_path),
+            force=False,
+            default_target=None,
+            unset_default_target=False,
+            external_access_integrations=["google_apis_access_integration", "dbt_hub"],
+        )
+
     def test_raises_when_dbt_project_yml_is_not_available(
         self, dbt_project_path, mock_connect, runner
     ):

tests/dbt/test_manager.py

@@ -48,10 +48,108 @@ def project_path(self, tmp_path_factory):
         source_path = tmp_path_factory.mktemp("dbt_project")
         yield source_path
 
+    @pytest.fixture
+    def dbt_project_path(self, project_path, profile):
+        dbt_project_file = project_path / "dbt_project.yml"
+        dbt_project_file.write_text(yaml.dump({"profile": "dev"}))
+        dbt_profiles_file = project_path / PROFILES_FILENAME
+        dbt_profiles_file.write_text(yaml.dump(profile))
+        yield project_path
+
     def _generate_profile(self, project_path, profile):
         dbt_profiles_file = project_path / PROFILES_FILENAME
         dbt_profiles_file.write_text(yaml.dump(profile))
 
+    @pytest.fixture
+    def mock_get_dbt_object_attributes(self):
+        with mock.patch(
+            "snowflake.cli._plugins.dbt.manager.DBTManager.get_dbt_object_attributes",
+            return_value=None,
+        ) as _fixture:
+            yield _fixture
+
+    @pytest.fixture
+    def mock_execute_query(self):
+        with mock.patch(
+            "snowflake.cli._plugins.dbt.manager.DBTManager.execute_query"
+        ) as _fixture:
+            yield _fixture
+
+    @pytest.fixture
+    def mock_get_cli_context(self, mock_connect):
+        with mock.patch(
+            "snowflake.cli.api.cli_global_context.get_cli_context"
+        ) as cli_context:
+            mock_connect.database = "TestDB"
+            mock_connect.schema = "TestSchema"
+            cli_context().connection = mock_connect
+            yield cli_context()
+
+    @pytest.fixture
+    def mock_from_resource(self):
+        with mock.patch(
+            "snowflake.cli._plugins.dbt.manager.FQN.from_resource",
+            return_value="@MockDatabase.MockSchema.DBT_PROJECT_TEST_PIPELINE_1757333281_STAGE",
+        ) as _fixture:
+            yield _fixture
+
+    @mock.patch("snowflake.cli._plugins.dbt.manager.StageManager.create")
+    @mock.patch("snowflake.cli._plugins.dbt.manager.StageManager.put_recursive")
+    def test_deploy_with_external_access_integrations(
+        self,
+        _mock_put_recursive,
+        _mock_create,
+        dbt_project_path,
+        mock_get_dbt_object_attributes,
+        mock_execute_query,
+        mock_get_cli_context,
+        mock_from_resource,
+    ):
+        manager = DBTManager()
+
+        manager.deploy(
+            fqn=FQN.from_string("test_project"),
+            path=SecurePath(dbt_project_path),
+            profiles_path=SecurePath(dbt_project_path),
+            force=False,
+            external_access_integrations=[
+                "google_apis_access_integration",
+                "dbt_hub_integration",
+            ],
+        )
+
+        expected_query = f"CREATE DBT PROJECT test_project\nFROM {mock_from_resource()}\nEXTERNAL_ACCESS_INTEGRATIONS = (google_apis_access_integration, dbt_hub_integration)"
+        mock_execute_query.assert_called_once_with(expected_query)
+
+    @mock.patch("snowflake.cli._plugins.dbt.manager.StageManager.create")
+    @mock.patch("snowflake.cli._plugins.dbt.manager.StageManager.put_recursive")
+    def test_deploy_alter_project_with_external_access_integrations(
+        self,
+        _mock_put_recursive,
+        _mock_create,
+        dbt_project_path,
+        mock_get_dbt_object_attributes,
+        mock_execute_query,
+        mock_get_cli_context,
+        mock_from_resource,
+    ):
+        mock_get_dbt_object_attributes.return_value = {"default_target": None}
+        manager = DBTManager()
+
+        manager.deploy(
+            fqn=FQN.from_string("test_project"),
+            path=SecurePath(dbt_project_path),
+            profiles_path=SecurePath(dbt_project_path),
+            force=False,
+            external_access_integrations=[
+                "google_apis_access_integration",
+                "dbt_hub_integration",
+            ],
+        )
+
+        expected_query = f"ALTER DBT PROJECT test_project ADD VERSION\nFROM {mock_from_resource()}\nEXTERNAL_ACCESS_INTEGRATIONS = (google_apis_access_integration, dbt_hub_integration)"
+        mock_execute_query.assert_called_once_with(expected_query)
+
     def test_validate_profiles_raises_when_file_does_not_exist(self, project_path):
 
         with pytest.raises(CliError) as exc_info:

tests_integration/test_data/projects/dbt_project_with_external_deps/dbt_project.yml

@@ -0,0 +1,20 @@
+name: 'dbt_integration_project'
+version: '1.0.0'
+
+profile: 'dbt_integration_project'
+
+model-paths: ["models"]
+analysis-paths: ["analyses"]
+test-paths: ["tests"]
+seed-paths: ["seeds"]
+macro-paths: ["macros"]
+snapshot-paths: ["snapshots"]
+
+clean-targets:
+  - "target"
+  - "dbt_packages"
+
+models:
+  dbt_integration_project:
+    example:
+      +materialized: view