Snow 2097403 add encryption to stage create (#2300)

* Add --encryption option to stage create

* fix unit tests

* integration tests

* update release notes

* review fixes

Author: sfc-gh-pczajka

RELEASE-NOTES.md

@@ -20,6 +20,7 @@
 
 ## New additions
 * Added `--private-link` flag to `snow spcs image-registry url` command for retrieving private link URLs.
+* Added `--encryption` flag to `snow stage create` command defining the type of encryption for all files on the stage.
 
 ## Fixes and improvements
 

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

@@ -29,7 +29,10 @@
     DiffResult,
     compute_stage_diff,
 )
-from snowflake.cli._plugins.stage.manager import StageManager
+from snowflake.cli._plugins.stage.manager import (
+    InternalStageEncryptionType,
+    StageManager,
+)
 from snowflake.cli._plugins.stage.utils import print_diff_to_console
 from snowflake.cli.api.cli_global_context import get_cli_context
 from snowflake.cli.api.commands.common import OnErrorType
@@ -150,11 +153,19 @@ def copy(
 
 
 @app.command("create", requires_connection=True)
-def stage_create(stage_name: FQN = StageNameArgument, **options) -> CommandResult:
+def stage_create(
+    stage_name: FQN = StageNameArgument,
+    encryption: InternalStageEncryptionType = typer.Option(
+        InternalStageEncryptionType.SNOWFLAKE_FULL.value,
+        "--encryption",
+        help="Type of encryption supported for all files stored on the stage.",
+    ),
+    **options,
+) -> CommandResult:
     """
     Creates a named stage if it does not already exist.
     """
-    cursor = StageManager().create(fqn=stage_name)
+    cursor = StageManager().create(fqn=stage_name, encryption=encryption)
     return SingleQueryResult(cursor)
 
 

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

@@ -25,6 +25,7 @@
 from collections import deque
 from contextlib import nullcontext
 from dataclasses import dataclass
+from enum import Enum
 from os import path
 from pathlib import Path
 from tempfile import TemporaryDirectory
@@ -68,6 +69,11 @@
 STAGE_PATH_REGEX = rf"(?P<prefix>(@|{re.escape('snow://')}))?(?:(?P<first_qualifier>{VALID_IDENTIFIER_REGEX})\.)?(?:(?P<second_qualifier>{VALID_IDENTIFIER_REGEX})\.)?(?P<name>{VALID_IDENTIFIER_REGEX})/?(?P<directory>([^/]*/?)*)?"
 
 
+class InternalStageEncryptionType(Enum):
+    SNOWFLAKE_FULL = "SNOWFLAKE_FULL"
+    SNOWFLAKE_SSE = "SNOWFLAKE_SSE"
+
+
 @dataclass
 class StagePathParts:
     directory: str
@@ -515,10 +521,16 @@ def remove(
             return self.execute_query(f"remove {stage_path.path_for_sql()}")
 
     def create(
-        self, fqn: FQN, comment: Optional[str] = None, temporary: bool = False
+        self,
+        fqn: FQN,
+        comment: Optional[str] = None,
+        temporary: bool = False,
+        encryption: InternalStageEncryptionType | None = None,
     ) -> SnowflakeCursor:
         temporary_str = "temporary " if temporary else ""
         query = f"create {temporary_str}stage if not exists {fqn.sql_identifier}"
+        if encryption:
+            query += f" encryption = (type = '{encryption.value}')"
         if comment:
             query += f" comment='{comment}'"
         return self.execute_query(query)

tests/__snapshots__/test_help_messages.ambr

@@ -15173,7 +15173,12 @@
   |                            [required]                                        |
   +------------------------------------------------------------------------------+
   +- Options --------------------------------------------------------------------+
-  | --help  -h        Show this message and exit.                                |
+  | --encryption          [SNOWFLAKE_FULL|SNOWFLAKE_  Type of encryption         |
+  |                       SSE]                        supported for all files    |
+  |                                                   stored on the stage.       |
+  |                                                   [default: SNOWFLAKE_FULL]  |
+  | --help        -h                                  Show this message and      |
+  |                                                   exit.                      |
   +------------------------------------------------------------------------------+
   +- Connection configuration ---------------------------------------------------+
   | --connection,--environment    -c      TEXT     Name of the connection, as    |

tests/stage/test_stage.py

@@ -114,7 +114,7 @@ def test_stage_copy_remote_to_local_quoted_stage_recursive(
 ):
     mock_execute.side_effect = [
         mock_cursor([{"name": '"stage name"/file'}], []),
-        mock_cursor([("file")], ["file"]),
+        mock_cursor(["file"], ["file"]),
     ]
     with TemporaryDirectory() as tmp_dir:
         result = runner.invoke(
@@ -191,7 +191,7 @@ def test_stage_copy_remote_to_local_quoted_uri_recursive(
 ):
     mock_execute.side_effect = [
         mock_cursor([{"name": "stageName/file.py"}], []),
-        mock_cursor([(raw_path)], ["file"]),
+        mock_cursor([raw_path], ["file"]),
     ]
     with TemporaryDirectory() as tmp_dir:
         tmp_dir = Path(tmp_dir).resolve()
@@ -516,17 +516,41 @@ def test_stage_create(mock_execute, runner, mock_cursor):
     result = runner.invoke(["stage", "create", "-c", "empty", "stageName"])
     assert result.exit_code == 0, result.output
     mock_execute.assert_called_once_with(
-        "create stage if not exists IDENTIFIER('stageName')"
+        "create stage if not exists IDENTIFIER('stageName') encryption = (type = 'SNOWFLAKE_FULL')"
     )
 
 
+@mock.patch(f"{STAGE_MANAGER}.execute_query")
+def test_stage_create_encryption(mock_execute, runner, mock_cursor):
+    mock_execute.return_value = mock_cursor(["row"], [])
+    for encryption in ["SNOWFLAKE_SSE", "SNOWFLAKE_FULL"]:
+        result = runner.invoke(
+            ["stage", "create", '"stage name"', "--encryption", encryption]
+        )
+        assert result.exit_code == 0, result.output
+        mock_execute.assert_called_once_with(
+            f"""create stage if not exists IDENTIFIER('"stage name"') encryption = (type = '{encryption}')"""
+        )
+        mock_execute.reset_mock()
+
+    result = runner.invoke(
+        ["stage", "create", '"stage name"', "--encryption", "incorrect_encryption"]
+    )
+    assert result.exit_code == 2, result.output
+    assert (
+        "Invalid value for '--encryption': 'incorrect_encryption' is not one of"
+        in result.output
+    )
+    assert "'SNOWFLAKE_FULL', 'SNOWFLAKE_SSE'." in result.output
+
+
 @mock.patch(f"{STAGE_MANAGER}.execute_query")
 def test_stage_create_quoted(mock_execute, runner, mock_cursor):
     mock_execute.return_value = mock_cursor(["row"], [])
     result = runner.invoke(["stage", "create", "-c", "empty", '"stage name"'])
     assert result.exit_code == 0, result.output
     mock_execute.assert_called_once_with(
-        """create stage if not exists IDENTIFIER('"stage name"')"""
+        """create stage if not exists IDENTIFIER('"stage name"') encryption = (type = 'SNOWFLAKE_FULL')"""
     )
 
 

tests_integration/test_stage.py

@@ -18,9 +18,7 @@
 import tempfile
 import time
 from pathlib import Path
-
 import pytest
-from snowflake.connector import DictCursor
 
 from tests.stage.test_stage import RecursiveUploadTester, NESTED_STRUCTURE
 from tests_integration.test_utils import (
@@ -819,3 +817,13 @@ def test_recursive_upload_no_recursive_glob_pattern(
             "target_size": 16,
         }
     ]
+
+
+@pytest.mark.integration
+@pytest.mark.parametrize("encryption", ["SNOWFLAKE_FULL", "SNOWFLAKE_SSE"])
+def test_create_encryption(runner, test_database, encryption):
+    result = runner.invoke_with_connection_json(
+        ["stage", "create", "a_stage", "--encryption", encryption]
+    )
+    assert result.exit_code == 0, result.output
+    assert result.json == {"status": f"Stage area A_STAGE successfully created."}