Snow 2097308 fix certificate issue on qa (#2322)

* remove snowflake.core from streamlit

* fix tests

Author: sfc-gh-pczajka

src/snowflake/cli/_plugins/streamlit/streamlit_entity.py

@@ -99,9 +99,6 @@ def deploy(
                 f"Streamlit {self.model.fqn.sql_identifier} already exists. Use 'replace' option to overwrite."
             )
 
-        console.step(f"Creating stage {self.model.stage} if not exists")
-        stage = self._create_stage_if_not_exists()
-
         if (
             experimental
             or GlobalFeatureFlag.ENABLE_STREAMLIT_VERSIONED_STAGE.is_enabled()
@@ -115,23 +112,20 @@ def deploy(
             name = (
                 self.model.identifier.name
                 if isinstance(self.model.identifier, Identifier)
-                else self.model.identifier
+                else self.model.identifier or self.entity_id
             )
             stage_root = StageManager.get_standard_stage_prefix(
                 f"{FQN.from_string(self.model.stage).using_connection(self._conn)}/{name}"
             )
-            if prune:
-                sync_deploy_root_with_stage(
-                    console=self._workspace_ctx.console,
-                    deploy_root=bundle_map.deploy_root(),
-                    bundle_map=bundle_map,
-                    prune=prune,
-                    recursive=True,
-                    stage_path=StageManager().stage_path_parts_from_str(stage_root),
-                    print_diff=True,
-                )
-            else:
-                self._upload_files_to_stage(stage, bundle_map, None)
+            sync_deploy_root_with_stage(
+                console=self._workspace_ctx.console,
+                deploy_root=bundle_map.deploy_root(),
+                bundle_map=bundle_map,
+                prune=prune,
+                recursive=True,
+                stage_path=StageManager().stage_path_parts_from_str(stage_root),
+                print_diff=True,
+            )
 
             console.step(f"Creating Streamlit object {self.model.fqn.sql_identifier}")
 
@@ -266,17 +260,12 @@ def _deploy_experimental(
         else:
             stage_root = f"{embeded_stage_name}/default_checkout"
 
-        stage_resource = self._create_stage_if_not_exists(embeded_stage_name)
-
-        if prune:
-            sync_deploy_root_with_stage(
-                console=self._workspace_ctx.console,
-                deploy_root=bundle_map.deploy_root(),
-                bundle_map=bundle_map,
-                prune=prune,
-                recursive=True,
-                stage_path=StageManager().stage_path_parts_from_str(stage_root),
-                print_diff=True,
-            )
-        else:
-            self._upload_files_to_stage(stage_resource, bundle_map)
+        sync_deploy_root_with_stage(
+            console=self._workspace_ctx.console,
+            deploy_root=bundle_map.deploy_root(),
+            bundle_map=bundle_map,
+            prune=prune,
+            recursive=True,
+            stage_path=StageManager().stage_path_parts_from_str(stage_root),
+            print_diff=True,
+        )

src/snowflake/cli/api/entities/common.py

@@ -4,18 +4,13 @@
 from typing import Generic, List, Optional, Type, TypeVar, get_args
 
 from snowflake.cli._plugins.workspace.context import ActionContext, WorkspaceContext
-from snowflake.cli.api.artifacts.bundle_map import BundleMap
 from snowflake.cli.api.cli_global_context import get_cli_context, span
 from snowflake.cli.api.entities.resolver import Dependency, DependencyResolver
 from snowflake.cli.api.entities.utils import EntityActions, get_sql_executor
 from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.sql_execution import SqlExecutor
-from snowflake.cli.api.utils.path_utils import change_directory
-from snowflake.cli.api.utils.python_api_utils import StageEncryptionType
 from snowflake.connector import SnowflakeConnection
 from snowflake.connector.cursor import SnowflakeCursor
-from snowflake.core import CreateMode
-from snowflake.core.stage import Stage, StageEncryption, StageResource
 
 T = TypeVar("T")
 
@@ -129,13 +124,6 @@ def snow_api_root(self) -> Optional[object]:
             raise ValueError("snow_api_root is not set")
         return root
 
-    @property
-    def stage_object(self) -> "StageResource":
-        if self._stage_object is None:
-            self._stage_object = self._create_stage_if_not_exists()
-
-        return self._stage_object
-
     @property
     def model(self) -> T:
         return self._entity_model
@@ -152,53 +140,13 @@ def get_describe_sql(self) -> str:
     def get_drop_sql(self) -> str:
         return f"DROP {self.model.type.upper()} {self.identifier};"  # type: ignore[attr-defined]
 
-    def _create_stage_if_not_exists(
-        self, stage_name: Optional[str] = None
-    ) -> StageResource:
-        if stage_name is None:
-            stage_name = self.model.stage  # type: ignore[attr-defined]
-
-        stage_collection = (
-            self.snow_api_root.databases[self.database].schemas[self.schema].stages  # type: ignore[attr-defined]
-        )
-        stage_object = Stage(
-            name=stage_name,
-            encryption=StageEncryption(type=StageEncryptionType.SNOWFLAKE_SSE.value),
-        )
-
-        return stage_collection.create(stage_object, mode=CreateMode.if_not_exists)
-
     def _get_identifier(
         self, schema: Optional[str] = None, database: Optional[str] = None
     ) -> str:
         schema_to_use = schema or self._entity_model.fqn.schema or self._conn.schema  # type: ignore
         db_to_use = database or self._entity_model.fqn.database or self._conn.database  # type: ignore
         return f"{self._entity_model.fqn.set_schema(schema_to_use).set_database(db_to_use).sql_identifier}"  # type: ignore
 
-    def _upload_files_to_stage(
-        self,
-        stage: StageResource,
-        bundle_map: BundleMap,
-        stage_root: Optional[str] = None,
-    ) -> None:
-        with change_directory(self.root):
-            for src, dest in bundle_map.all_mappings(
-                absolute=True, expand_directories=True
-            ):
-                if src.is_file():
-                    upload_dst = (
-                        f"{stage_root}/{dest.relative_to(self.root)}"
-                        if stage_root
-                        else f"/{self.fqn.name}/{get_parent_path_for_stage_deployment(dest.relative_to(bundle_map.deploy_root()))}"
-                    )
-
-                    stage.put(
-                        local_file_name=src.relative_to(self.root),
-                        stage_location=upload_dst,
-                        overwrite=True,
-                        auto_compress=False,
-                    )
-
     def get_from_fqn_or_conn(self, attribute_name: str) -> str:
         attribute = getattr(self.fqn, attribute_name, None) or getattr(
             self._conn, attribute_name, None

tests/conftest.py

@@ -334,39 +334,41 @@ def execute_stream(self, query: StringIO, **kwargs):
         return self.execute_string(query.read(), **kwargs)
 
 
-@pytest.fixture
-def mock_cursor():
-    class MockResultMetadata(NamedTuple):
-        name: str
+class MockResultMetadata(NamedTuple):
+    name: str
+
+
+class MockCursor(SnowflakeCursor):
+    def __init__(self, rows: List[Union[tuple, dict]], columns: List[str]):
+        super().__init__(mock.Mock())
+        self._rows = rows
+        self._columns = [MockResultMetadata(c) for c in columns]
+        self.query = "SELECT A MOCK QUERY"
 
-    class _MockCursor(SnowflakeCursor):
-        def __init__(self, rows: List[Union[tuple, dict]], columns: List[str]):
-            super().__init__(mock.Mock())
-            self._rows = rows
-            self._columns = [MockResultMetadata(c) for c in columns]
-            self.query = "SELECT A MOCK QUERY"
+    def fetchone(self):
+        if self._rows:
+            return self._rows.pop(0)
+        return None
 
-        def fetchone(self):
-            if self._rows:
-                return self._rows.pop(0)
-            return None
+    def fetchall(self):
+        return self._rows
 
-        def fetchall(self):
-            return self._rows
+    @property
+    def rowcount(self):
+        return len(self._rows)
 
-        @property
-        def rowcount(self):
-            return len(self._rows)
+    @property
+    def description(self):
+        yield from self._columns
 
-        @property
-        def description(self):
-            yield from self._columns
+    @classmethod
+    def from_input(cls, rows, columns):
+        return cls(rows, columns)
 
-        @classmethod
-        def from_input(cls, rows, columns):
-            return cls(rows, columns)
 
-    return _MockCursor.from_input
+@pytest.fixture
+def mock_cursor():
+    return MockCursor.from_input
 
 
 @pytest.fixture
@@ -512,12 +514,15 @@ def native_app_project_instance():
 
 @pytest.fixture
 def enable_snowpark_glob_support_feature_flag():
-    with mock.patch(
-        f"snowflake.cli.api.feature_flags.FeatureFlag.ENABLE_SNOWPARK_GLOB_SUPPORT.is_enabled",
-        return_value=True,
-    ), mock.patch(
-        f"snowflake.cli.api.feature_flags.FeatureFlag.ENABLE_SNOWPARK_GLOB_SUPPORT.is_disabled",
-        return_value=False,
+    with (
+        mock.patch(
+            f"snowflake.cli.api.feature_flags.FeatureFlag.ENABLE_SNOWPARK_GLOB_SUPPORT.is_enabled",
+            return_value=True,
+        ),
+        mock.patch(
+            f"snowflake.cli.api.feature_flags.FeatureFlag.ENABLE_SNOWPARK_GLOB_SUPPORT.is_disabled",
+            return_value=False,
+        ),
     ):
         yield
 

tests/streamlit/streamlit_test_class.py

@@ -1,9 +1,8 @@
-from pathlib import Path
+from re import match
 from typing import List
 from unittest import mock
-from unittest.mock import MagicMock
 
-from snowflake.core.stage import StageResource
+from tests.conftest import MockCursor
 
 STREAMLIT_NAME = "test_streamlit"
 
@@ -26,16 +25,17 @@ def setup_method(self):
 
         self.mock_execute = mock.patch(EXECUTE_QUERY).start()
 
-        mock_stage_resource = StageResource(
-            name="stage_resource_mock", collection=MagicMock()
-        )
         self.mock_create_stage = mock.patch(
-            "snowflake.cli._plugins.streamlit.streamlit_entity.StreamlitEntity._create_stage_if_not_exists",
-            return_value=mock_stage_resource,
+            "snowflake.cli._plugins.stage.manager.StageManager.create",
+        ).start()
+
+        self.mock_list_files = mock.patch(
+            "snowflake.cli._plugins.stage.manager.StageManager.list_files",
+            return_value=MockCursor.from_input([], []),
         ).start()
 
         self.mock_put = mock.patch(
-            "snowflake.core.stage._stage.StageResource.put"
+            "snowflake.cli._plugins.stage.manager.StageManager.put"
         ).start()
 
         self.mock_get_account = mock.patch(
@@ -61,19 +61,14 @@ def _assert_that_exactly_those_files_were_put_to_stage(
         put_files: List[str],
         streamlit_name: str = "test_streamlit",
     ):
-        assert self.mock_put.call_count == len(put_files)  # type: ignore
-
-        for file in put_files:
-            if isinstance(file, dict):
-                local = Path(file["local"])
-                stage = f"/{streamlit_name}{file['stage'] if file['stage'] else ''}"
-            else:
-                local = Path(file)
-                stage = f"/{streamlit_name}/{str(Path(file).parent) if Path(file).parent != Path('.') else ''}"
-
-            self.mock_put.assert_any_call(  # type: ignore
-                local_file_name=local,
-                stage_location=stage,
-                overwrite=True,
-                auto_compress=False,
-            )
+        # assert self.mock_put.call_count == len(put_files)  # type: ignore
+
+        re_local_path = f".*/{streamlit_name}/(?P<filename>.*)"
+        uploaded_files = set()
+        for call in self.mock_put.call_args_list:
+            if path := call.kwargs.get("local_path"):
+                matched_path = match(re_local_path, path.as_posix())
+                if matched_path:
+                    uploaded_files.add(matched_path.group("filename"))
+
+        assert set(put_files) == uploaded_files, uploaded_files