Workspaces application package - validate setup script (#1495)

* extract validate logic to entity * ws deploy validate option * assert validation
snowflakedb · Aug 27, 2024 · 110dcb8 · 110dcb8
1 parent 90cce15
commit 110dcb8
Show file tree

Hide file tree

Showing 9 changed files with 191 additions and 84 deletions.
diff --git a/src/snowflake/cli/_plugins/nativeapp/manager.py b/src/snowflake/cli/_plugins/nativeapp/manager.py
@@ -14,7 +14,6 @@
 
 from __future__ import annotations
 
-import json
 import time
 from abc import ABC, abstractmethod
 from contextlib import contextmanager
@@ -37,17 +36,14 @@
     NAME_COL,
 )
 from snowflake.cli._plugins.nativeapp.exceptions import (
-    ApplicationPackageDoesNotExistError,
     NoEventTableForAccount,
-    SetupScriptFailedValidation,
 )
 from snowflake.cli._plugins.nativeapp.project_model import (
     NativeAppProjectModel,
 )
 from snowflake.cli._plugins.stage.diff import (
     DiffResult,
 )
-from snowflake.cli._plugins.stage.manager import StageManager
 from snowflake.cli.api.console import cli_console as cc
 from snowflake.cli.api.entities.application_package_entity import (
     ApplicationPackageEntity,
@@ -57,10 +53,6 @@
     generic_sql_error_handler,
     sync_deploy_root_with_stage,
 )
-from snowflake.cli.api.errno import (
-    DOES_NOT_EXIST_OR_NOT_AUTHORIZED,
-)
-from snowflake.cli.api.exceptions import SnowflakeSQLExecutionError
 from snowflake.cli.api.project.schemas.entities.common import PostDeployHook
 from snowflake.cli.api.project.schemas.native_app.native_app import NativeApp
 from snowflake.cli.api.project.schemas.native_app.path_mapping import PathMapping
@@ -381,59 +373,38 @@ def deploy(
 
         return diff
 
-    def validate(self, use_scratch_stage: bool = False):
-        """Validates Native App setup script SQL."""
-        with cc.phase(f"Validating Snowflake Native App setup script."):
-            validation_result = self.get_validation_result(use_scratch_stage)
-
-            # First print warnings, regardless of the outcome of validation
-            for warning in validation_result.get("warnings", []):
-                cc.warning(_validation_item_to_str(warning))
-
-            # Then print errors
-            for error in validation_result.get("errors", []):
-                # Print them as warnings for now since we're going to be
-                # revamping CLI output soon
-                cc.warning(_validation_item_to_str(error))
+    def deploy_to_scratch_stage_fn(self):
+        bundle_map = self.build_bundle()
+        self.deploy(
+            bundle_map=bundle_map,
+            prune=True,
+            recursive=True,
+            stage_fqn=self.scratch_stage_fqn,
+            validate=False,
+            print_diff=False,
+        )
 
-            # Then raise an exception if validation failed
-            if validation_result["status"] == "FAIL":
-                raise SetupScriptFailedValidation()
+    def validate(self, use_scratch_stage: bool = False):
+        return ApplicationPackageEntity.validate_setup_script(
+            console=cc,
+            package_name=self.package_name,
+            package_role=self.package_role,
+            stage_fqn=self.stage_fqn,
+            use_scratch_stage=use_scratch_stage,
+            scratch_stage_fqn=self.scratch_stage_fqn,
+            deploy_to_scratch_stage_fn=self.deploy_to_scratch_stage_fn,
+        )
 
     def get_validation_result(self, use_scratch_stage: bool):
-        """Call system$validate_native_app_setup() to validate deployed Native App setup script."""
-        stage_fqn = self.stage_fqn
-        if use_scratch_stage:
-            stage_fqn = self.scratch_stage_fqn
-            bundle_map = self.build_bundle()
-            self.deploy(
-                bundle_map=bundle_map,
-                prune=True,
-                recursive=True,
-                stage_fqn=stage_fqn,
-                validate=False,
-                print_diff=False,
-            )
-        prefixed_stage_fqn = StageManager.get_standard_stage_prefix(stage_fqn)
-        try:
-            cursor = self._execute_query(
-                f"call system$validate_native_app_setup('{prefixed_stage_fqn}')"
-            )
-        except ProgrammingError as err:
-            if err.errno == DOES_NOT_EXIST_OR_NOT_AUTHORIZED:
-                raise ApplicationPackageDoesNotExistError(self.package_name)
-            generic_sql_error_handler(err)
-        else:
-            if not cursor.rowcount:
-                raise SnowflakeSQLExecutionError()
-            return json.loads(cursor.fetchone()[0])
-        finally:
-            if use_scratch_stage:
-                cc.step(f"Dropping stage {self.scratch_stage_fqn}.")
-                with self.use_role(self.package_role):
-                    self._execute_query(
-                        f"drop stage if exists {self.scratch_stage_fqn}"
-                    )
+        return ApplicationPackageEntity.get_validation_result(
+            console=cc,
+            package_name=self.package_name,
+            package_role=self.package_role,
+            stage_fqn=self.stage_fqn,
+            use_scratch_stage=use_scratch_stage,
+            scratch_stage_fqn=self.scratch_stage_fqn,
+            deploy_to_scratch_stage_fn=self.deploy_to_scratch_stage_fn,
+        )
 
     def get_events(  # type: ignore [return]
         self,

diff --git a/src/snowflake/cli/_plugins/workspace/commands.py b/src/snowflake/cli/_plugins/workspace/commands.py
@@ -23,6 +23,7 @@
 import yaml
 from click import ClickException
 from snowflake.cli._plugins.nativeapp.artifacts import BundleMap
+from snowflake.cli._plugins.nativeapp.common_flags import ValidateOption
 from snowflake.cli._plugins.snowpark.commands import migrate_v1_snowpark_to_v2
 from snowflake.cli._plugins.streamlit.commands import migrate_v1_streamlit_to_v2
 from snowflake.cli._plugins.workspace.manager import WorkspaceManager
@@ -145,6 +146,7 @@ def deploy(
             unspecified, the command syncs all local changes to the stage."""
         ).strip(),
     ),
+    validate: bool = ValidateOption,
     **options,
 ):
     """
@@ -174,5 +176,6 @@ def deploy(
         prune=prune,
         recursive=recursive,
         paths=paths,
+        validate=validate,
     )
     return MessageResult("Deployed successfully.")
diff --git a/src/snowflake/cli/api/entities/application_package_entity.py b/src/snowflake/cli/api/entities/application_package_entity.py
@@ -1,7 +1,8 @@
+import json
 from contextlib import contextmanager
 from pathlib import Path
 from textwrap import dedent
-from typing import List, Optional
+from typing import Callable, List, Optional
 
 from click import ClickException
 from snowflake.cli._plugins.nativeapp.artifacts import build_bundle
@@ -16,7 +17,10 @@
 )
 from snowflake.cli._plugins.nativeapp.exceptions import (
     ApplicationPackageAlreadyExistsError,
+    ApplicationPackageDoesNotExistError,
+    SetupScriptFailedValidation,
 )
+from snowflake.cli._plugins.stage.manager import StageManager
 from snowflake.cli._plugins.workspace.action_context import ActionContext
 from snowflake.cli.api.console.abc import AbstractConsole
 from snowflake.cli.api.entities.common import EntityBase, get_sql_executor
@@ -25,6 +29,10 @@
     generic_sql_error_handler,
     render_script_templates,
     sync_deploy_root_with_stage,
+    validation_item_to_str,
+)
+from snowflake.cli.api.errno import (
+    DOES_NOT_EXIST_OR_NOT_AUTHORIZED,
 )
 from snowflake.cli.api.exceptions import SnowflakeSQLExecutionError
 from snowflake.cli.api.project.schemas.entities.application_package_entity_model import (
@@ -65,6 +73,7 @@ def action_deploy(
         prune: bool,
         recursive: bool,
         paths: List[Path],
+        validate: bool,
     ):
         model = self._entity_model
         package_name = model.fqn.identifier
@@ -103,7 +112,17 @@ def action_deploy(
             )
 
         # TODO Execute post-deploy hooks
-        # TODO Validate
+
+        if validate:
+            self.validate_setup_script(
+                console=ctx.console,
+                package_name=package_name,
+                package_role=package_role,
+                stage_fqn=stage_fqn,
+                use_scratch_stage=False,
+                scratch_stage_fqn="",
+                deploy_to_scratch_stage_fn=lambda *args: None,
+            )
 
     @staticmethod
     def get_existing_app_pkg_info(
@@ -306,3 +325,76 @@ def create_app_package(
                 """
                 )
             )
+
+    @classmethod
+    def validate_setup_script(
+        cls,
+        console: AbstractConsole,
+        package_name: str,
+        package_role: str,
+        stage_fqn: str,
+        use_scratch_stage: bool,
+        scratch_stage_fqn: str,
+        deploy_to_scratch_stage_fn: Callable,
+    ):
+        """Validates Native App setup script SQL."""
+        with console.phase(f"Validating Snowflake Native App setup script."):
+            validation_result = cls.get_validation_result(
+                console=console,
+                package_name=package_name,
+                package_role=package_role,
+                stage_fqn=stage_fqn,
+                use_scratch_stage=use_scratch_stage,
+                scratch_stage_fqn=scratch_stage_fqn,
+                deploy_to_scratch_stage_fn=deploy_to_scratch_stage_fn,
+            )
+
+            # First print warnings, regardless of the outcome of validation
+            for warning in validation_result.get("warnings", []):
+                console.warning(validation_item_to_str(warning))
+
+            # Then print errors
+            for error in validation_result.get("errors", []):
+                # Print them as warnings for now since we're going to be
+                # revamping CLI output soon
+                console.warning(validation_item_to_str(error))
+
+            # Then raise an exception if validation failed
+            if validation_result["status"] == "FAIL":
+                raise SetupScriptFailedValidation()
+
+    @staticmethod
+    def get_validation_result(
+        console: AbstractConsole,
+        package_name: str,
+        package_role: str,
+        stage_fqn: str,
+        use_scratch_stage: bool,
+        scratch_stage_fqn: str,
+        deploy_to_scratch_stage_fn: Callable,
+    ):
+        """Call system$validate_native_app_setup() to validate deployed Native App setup script."""
+        if use_scratch_stage:
+            stage_fqn = scratch_stage_fqn
+            deploy_to_scratch_stage_fn()
+        prefixed_stage_fqn = StageManager.get_standard_stage_prefix(stage_fqn)
+        sql_executor = get_sql_executor()
+        try:
+            cursor = sql_executor.execute_query(
+                f"call system$validate_native_app_setup('{prefixed_stage_fqn}')"
+            )
+        except ProgrammingError as err:
+            if err.errno == DOES_NOT_EXIST_OR_NOT_AUTHORIZED:
+                raise ApplicationPackageDoesNotExistError(package_name)
+            generic_sql_error_handler(err)
+        else:
+            if not cursor.rowcount:
+                raise SnowflakeSQLExecutionError()
+            return json.loads(cursor.fetchone()[0])
+        finally:
+            if use_scratch_stage:
+                console.step(f"Dropping stage {scratch_stage_fqn}.")
+                with sql_executor.use_role(package_role):
+                    sql_executor.execute_query(
+                        f"drop stage if exists {scratch_stage_fqn}"
+                    )
diff --git a/src/snowflake/cli/api/entities/utils.py b/src/snowflake/cli/api/entities/utils.py
@@ -319,3 +319,10 @@ def render_script_templates(
             raise InvalidScriptError(relpath, e) from e
 
     return scripts_contents
+
+
+def validation_item_to_str(item: dict[str, str | int]):
+    s = item["message"]
+    if item["errorCode"]:
+        s = f"{s} (error code {item['errorCode']})"
+    return s
diff --git a/tests/nativeapp/test_manager.py b/tests/nativeapp/test_manager.py
@@ -40,7 +40,6 @@
 )
 from snowflake.cli._plugins.nativeapp.manager import (
     NativeAppManager,
-    SnowflakeSQLExecutionError,
 )
 from snowflake.cli._plugins.stage.diff import (
     DiffResult,
@@ -51,6 +50,7 @@
     ensure_correct_owner,
 )
 from snowflake.cli.api.errno import DOES_NOT_EXIST_OR_NOT_AUTHORIZED
+from snowflake.cli.api.exceptions import SnowflakeSQLExecutionError
 from snowflake.cli.api.project.definition_manager import DefinitionManager
 from snowflake.connector import ProgrammingError
 from snowflake.connector.cursor import DictCursor
@@ -957,7 +957,7 @@ def test_get_paths_to_sync(
     assert result.sort() == [StagePath(p) for p in expected_result].sort()
 
 
-@mock.patch(NATIVEAPP_MANAGER_EXECUTE)
+@mock.patch(SQL_EXECUTOR_EXECUTE)
 def test_validate_passing(mock_execute, temp_dir, mock_cursor):
     create_named_file(
         file_name="snowflake.yml",
@@ -984,7 +984,7 @@ def test_validate_passing(mock_execute, temp_dir, mock_cursor):
     assert mock_execute.mock_calls == expected
 
 
-@mock.patch(NATIVEAPP_MANAGER_EXECUTE)
+@mock.patch(SQL_EXECUTOR_EXECUTE)
 @mock.patch(f"{NATIVEAPP_MODULE}.cc.warning")
 def test_validate_passing_with_warnings(
     mock_warning, mock_execute, temp_dir, mock_cursor
@@ -1026,7 +1026,7 @@ def test_validate_passing_with_warnings(
     assert mock_execute.mock_calls == expected
 
 
-@mock.patch(NATIVEAPP_MANAGER_EXECUTE)
+@mock.patch(SQL_EXECUTOR_EXECUTE)
 @mock.patch(f"{NATIVEAPP_MODULE}.cc.warning")
 def test_validate_failing(mock_warning, mock_execute, temp_dir, mock_cursor):
     create_named_file(
@@ -1083,7 +1083,7 @@ def test_validate_failing(mock_warning, mock_execute, temp_dir, mock_cursor):
     assert mock_execute.mock_calls == expected
 
 
-@mock.patch(NATIVEAPP_MANAGER_EXECUTE)
+@mock.patch(SQL_EXECUTOR_EXECUTE)
 def test_validate_query_error(mock_execute, temp_dir, mock_cursor):
     create_named_file(
         file_name="snowflake.yml",
@@ -1110,7 +1110,7 @@ def test_validate_query_error(mock_execute, temp_dir, mock_cursor):
     assert mock_execute.mock_calls == expected
 
 
-@mock.patch(NATIVEAPP_MANAGER_EXECUTE)
+@mock.patch(SQL_EXECUTOR_EXECUTE)
 def test_validate_not_deployed(mock_execute, temp_dir, mock_cursor):
     create_named_file(
         file_name="snowflake.yml",
@@ -1142,7 +1142,7 @@ def test_validate_not_deployed(mock_execute, temp_dir, mock_cursor):
 
 @mock.patch(NATIVEAPP_MANAGER_BUILD_BUNDLE)
 @mock.patch(NATIVEAPP_MANAGER_DEPLOY)
-@mock.patch(NATIVEAPP_MANAGER_EXECUTE)
+@mock.patch(SQL_EXECUTOR_EXECUTE)
 def test_validate_use_scratch_stage(
     mock_execute, mock_deploy, mock_build_bundle, temp_dir, mock_cursor
 ):
@@ -1194,7 +1194,7 @@ def test_validate_use_scratch_stage(
 
 @mock.patch(NATIVEAPP_MANAGER_BUILD_BUNDLE)
 @mock.patch(NATIVEAPP_MANAGER_DEPLOY)
-@mock.patch(NATIVEAPP_MANAGER_EXECUTE)
+@mock.patch(SQL_EXECUTOR_EXECUTE)
 def test_validate_failing_drops_scratch_stage(
     mock_execute, mock_deploy, mock_build_bundle, temp_dir, mock_cursor
 ):
@@ -1258,7 +1258,7 @@ def test_validate_failing_drops_scratch_stage(
     assert mock_execute.mock_calls == expected
 
 
-@mock.patch(NATIVEAPP_MANAGER_EXECUTE)
+@mock.patch(SQL_EXECUTOR_EXECUTE)
 def test_validate_raw_returns_data(mock_execute, temp_dir, mock_cursor):
     create_named_file(
         file_name="snowflake.yml",

diff --git a/tests/nativeapp/test_teardown_processor.py b/tests/nativeapp/test_teardown_processor.py
@@ -26,14 +26,14 @@
     CouldNotDropApplicationPackageWithVersions,
     UnexpectedOwnerError,
 )
-from snowflake.cli._plugins.nativeapp.manager import SnowflakeSQLExecutionError
 from snowflake.cli._plugins.nativeapp.teardown_processor import (
     NativeAppTeardownProcessor,
 )
 from snowflake.cli.api.errno import (
     APPLICATION_NO_LONGER_AVAILABLE,
     DOES_NOT_EXIST_OR_CANNOT_BE_PERFORMED,
 )
+from snowflake.cli.api.exceptions import SnowflakeSQLExecutionError
 from snowflake.cli.api.project.definition_manager import DefinitionManager
 from snowflake.connector import ProgrammingError
 from snowflake.connector.cursor import DictCursor