From 6ccfdc8d08f3bf8999338073209c2dfd6509a43f Mon Sep 17 00:00:00 2001 From: Jan Sikorski Date: Wed, 11 Dec 2024 13:47:28 +0100 Subject: [PATCH] fix --- .../_plugins/streamlit/streamlit_entity.py | 65 +++++++++++++++---- .../__snapshots__/test_commands.ambr | 45 ++----------- tests/streamlit/test_actions.py | 30 +++++++-- .../example_streamlit_v2/snowflake.yml | 8 ++- 4 files changed, 86 insertions(+), 62 deletions(-) diff --git a/src/snowflake/cli/_plugins/streamlit/streamlit_entity.py b/src/snowflake/cli/_plugins/streamlit/streamlit_entity.py index 8f4e32fcc1..98ca31c09d 100644 --- a/src/snowflake/cli/_plugins/streamlit/streamlit_entity.py +++ b/src/snowflake/cli/_plugins/streamlit/streamlit_entity.py @@ -1,5 +1,6 @@ import functools -import shutil + +from typing import Optional from snowflake.cli._plugins.connection.util import make_snowsight_url from snowflake.cli._plugins.streamlit.streamlit_entity_model import ( @@ -9,6 +10,8 @@ from snowflake.cli.api.entities.common import EntityBase, get_sql_executor from snowflake.connector.cursor import SnowflakeCursor +from snowflake.cli.api.secure_path import SecurePath + class StreamlitEntity(EntityBase[StreamlitEntityModel]): """ @@ -31,6 +34,10 @@ def _sql_executor(self): def _conn(self): return self._sql_executor._conn # noqa + @property + def model(self): + return self._entity_model # noqa + def action_bundle(self, ctx: ActionContext, *args, **kwargs): # get all files from the model artifacts = self._entity_model.artifacts @@ -45,10 +52,10 @@ def action_bundle(self, ctx: ActionContext, *args, **kwargs): output_file = output_folder / file.name if file.is_file(): - shutil.copy(file, output_file) + SecurePath(file).copy(output_file) elif file.is_dir(): output_file.mkdir(parents=True, exist_ok=True) - shutil.copytree(file, output_file, dirs_exist_ok=True) + SecurePath(file).copy(output_file, dirs_exist_ok=True) output_files.append(output_file) @@ -57,17 +64,17 @@ def action_bundle(self, ctx: ActionContext, *args, **kwargs): def action_deploy(self, action_ctx: ActionContext, *args, **kwargs): # After adding bundle map- we should use it's mapping here - query = self.action_get_deploy_sql(action_ctx, *args, **kwargs) + query = self.get_deploy_sql(action_ctx, *args, **kwargs) result = self._sql_executor.execute_query(query) return result def action_drop(self, action_ctx: ActionContext, *args, **kwargs): - return self._sql_executor.execute_query(self.action_get_drop_sql(action_ctx)) + return self._sql_executor.execute_query(self.get_drop_sql(action_ctx)) def action_execute( self, action_ctx: ActionContext, *args, **kwargs ) -> SnowflakeCursor: - return self._sql_executor.execute_query(self.action_get_execute_sql(action_ctx)) + return self._sql_executor.execute_query(self.get_execute_sql(action_ctx)) def action_get_url( self, action_ctx: ActionContext, *args, **kwargs @@ -77,21 +84,55 @@ def action_get_url( self._conn, f"/#/streamlit-apps/{name.url_identifier}" ) - def action_get_deploy_sql(self, action_ctx: ActionContext, *args, **kwargs): - pass + def get_deploy_sql(self, action_ctx: ActionContext, if_not_exists: bool = False, replace: bool = False, from_stage_name: Optional[str] = None, *args, **kwargs): + + if replace: + query = "CREATE OR REPLACE " + elif if_not_exists: + query = "CREATE IF NOT EXISTS " + else: + query = "CREATE " + + query += f"STREAMLIT {self._entity_model.fqn.sql_identifier} \n" + + if from_stage_name: + query += f" ROOT_LOCATION = '{from_stage_name}' \n" + + query += f" MAIN_FILE = '{self._entity_model.main_file}' \n" + + if self.model.imports: + query += self.model.get_imports_sql() + "\n" + + if self.model.query_warehouse: + query += f" QUERY_WAREHOUSE = '{self.model.query_warehouse}' \n" + + if self.model.title: + query += f" TITLE = '{self.model.title}' \n" + + if self.model.comment: + query += f" COMMENT = '{self.model.comment}' \n" + + if self.model.external_access_integrations: + query += self.model.get_external_access_integrations_sql() + "\n" + + if self.model.secrets: + query += self.model.get_secrets_sql() + "\n" + + return query + def action_share( self, action_ctx: ActionContext, to_role: str, *args, **kwargs ) -> SnowflakeCursor: - return self._sql_executor.execute_query(self.get_share_sql(action_ctx, to_role)) + return self._sql_executor.execute_query(self.get_usage_grant_sql(action_ctx, to_role)) - def action_get_drop_sql(self, action_ctx: ActionContext, *args, **kwargs): + def get_drop_sql(self, action_ctx: ActionContext, *args, **kwargs): return f"DROP STREAMLIT {self._entity_model.fqn}" - def action_get_execute_sql(self, action_ctx: ActionContext, *args, **kwargs): + def get_execute_sql(self, action_ctx: ActionContext, *args, **kwargs): return f"EXECUTE STREAMLIT {self._entity_model.fqn}()" - def get_share_sql( + def get_usage_grant_sql( self, action_ctx: ActionContext, to_role: str, *args, **kwargs ) -> str: return f"GRANT USAGE ON STREAMLIT {{self._entity_model.fqn}} to role {to_role}" diff --git a/tests/streamlit/__snapshots__/test_commands.ambr b/tests/streamlit/__snapshots__/test_commands.ambr index 846562546d..d86c8a9528 100644 --- a/tests/streamlit/__snapshots__/test_commands.ambr +++ b/tests/streamlit/__snapshots__/test_commands.ambr @@ -7,51 +7,16 @@ | For field entities.my_streamlit you provided '{'artifacts': {'1': | | 'foo_bar.py'}}'. This caused: Unable to extract tag using discriminator | | 'type' | - | For field entities.my_streamlit you provided '{'artifacts': {'1': | - | 'foo_bar.py'}}'. This caused: Unable to extract tag using discriminator | - | 'type' | +------------------------------------------------------------------------------+ - + ''' # --- -# name: test_deploy_put_files_on_stage[example_streamlit-merge_definition1] - list([ - "create stage if not exists IDENTIFIER('MockDatabase.MockSchema.streamlit_stage')", - 'put file://streamlit_app.py @MockDatabase.MockSchema.streamlit_stage/test_streamlit auto_compress=false parallel=4 overwrite=True', - 'put file://environment.yml @MockDatabase.MockSchema.streamlit_stage/test_streamlit auto_compress=false parallel=4 overwrite=True', - 'put file://pages/* @MockDatabase.MockSchema.streamlit_stage/test_streamlit/pages auto_compress=false parallel=4 overwrite=True', - ''' - CREATE STREAMLIT IDENTIFIER('MockDatabase.MockSchema.test_streamlit') - ROOT_LOCATION = '@MockDatabase.MockSchema.streamlit_stage/test_streamlit' - MAIN_FILE = 'streamlit_app.py' - QUERY_WAREHOUSE = test_warehouse - TITLE = 'My Fancy Streamlit' - ''', - 'select system$get_snowsight_host()', - 'select current_account_name()', - ]) -# --- -# name: test_deploy_put_files_on_stage[example_streamlit_v2-merge_definition0] - list([ - "create stage if not exists IDENTIFIER('MockDatabase.MockSchema.streamlit_stage')", - 'put file://streamlit_app.py @MockDatabase.MockSchema.streamlit_stage/test_streamlit auto_compress=false parallel=4 overwrite=True', - ''' - CREATE STREAMLIT IDENTIFIER('MockDatabase.MockSchema.test_streamlit') - ROOT_LOCATION = '@MockDatabase.MockSchema.streamlit_stage/test_streamlit' - MAIN_FILE = 'streamlit_app.py' - QUERY_WAREHOUSE = test_warehouse - TITLE = 'My Fancy Streamlit' - ''', - 'select system$get_snowsight_host()', - 'select current_account_name()', - ]) -# --- # name: test_deploy_streamlit_nonexisting_file[example_streamlit-opts0] ''' +- Error ----------------------------------------------------------------------+ | Provided file foo.bar does not exist | +------------------------------------------------------------------------------+ - + ''' # --- # name: test_deploy_streamlit_nonexisting_file[example_streamlit-opts1] @@ -59,7 +24,7 @@ +- Error ----------------------------------------------------------------------+ | Provided file foo.bar does not exist | +------------------------------------------------------------------------------+ - + ''' # --- # name: test_deploy_streamlit_nonexisting_file[example_streamlit_v2-opts2] @@ -68,7 +33,7 @@ | Streamlit test_streamlit already exist. If you want to replace it use | | --replace flag. | +------------------------------------------------------------------------------+ - + ''' # --- # name: test_deploy_streamlit_nonexisting_file[example_streamlit_v2-opts3] @@ -77,7 +42,7 @@ | During evaluation of DefinitionV20 in project definition following errors | | were encountered: | | For field entities.test_streamlit.streamlit you provided '{'artifacts': | - | ['foo.bar'], 'identifier': {'name': 'test_streamlit'}, 'main_file': | + | ['foo.bar'], 'identifier': 'test_streamlit', 'main_file': | | 'streamlit_app.py', 'query_warehouse': 'test_warehouse', 'stage': | | 'streamlit', 'title': 'My Fancy Streamlit', 'type': 'streamlit'}'. This | | caused: Value error, Specified artifact foo.bar does not exist locally. | diff --git a/tests/streamlit/test_actions.py b/tests/streamlit/test_actions.py index 0c870fa162..6a878f1904 100644 --- a/tests/streamlit/test_actions.py +++ b/tests/streamlit/test_actions.py @@ -7,6 +7,7 @@ from snowflake.cli._plugins.streamlit.streamlit_entity_model import StreamlitEntityModel from snowflake.cli._plugins.workspace.context import ActionContext, WorkspaceContext +STREAMLIT_NAME = "test_streamlit" CONNECTOR = "snowflake.connector.connect" CONTEXT = "" EXECUTE_QUERY = "snowflake.cli.api.sql_execution.BaseSqlExecutor.execute_query" @@ -19,7 +20,6 @@ def example_streamlit_workspace(project_directory): with project_directory("example_streamlit_v2") as pdir: with Path(pdir / "snowflake.yml").open() as definition_file: definition = yaml.safe_load(definition_file) - print(definition) model = StreamlitEntityModel( **definition.get("entities", {}).get("test_streamlit") ) @@ -50,13 +50,19 @@ def test_bundle(example_streamlit_workspace): assert (output / "environment.yml").exists() assert (output / "pages" / "my_page.py").exists() +@mock.patch(EXECUTE_QUERY) +def test_deploy(mock_execute, example_streamlit_workspace): + entity, action_ctx = example_streamlit_workspace + entity.action_deploy(action_ctx) + + mock_execute.assert_called_with(f"CREATE STREAMLIT IDENTIFIER('{STREAMLIT_NAME}') \n MAIN_FILE = 'streamlit_app.py' \n QUERY_WAREHOUSE = 'test_warehouse' \n TITLE = 'My Fancy Streamlit' \n") @mock.patch(EXECUTE_QUERY) def test_drop(mock_execute, example_streamlit_workspace): entity, action_ctx = example_streamlit_workspace entity.action_drop(action_ctx) - mock_execute.assert_called_with("DROP STREAMLIT test_streamlit_deploy_snowcli") + mock_execute.assert_called_with(f"DROP STREAMLIT {STREAMLIT_NAME}") @mock.patch(CONNECTOR) @@ -97,18 +103,28 @@ def test_execute(mock_execute, example_streamlit_workspace): entity, action_ctx = example_streamlit_workspace entity.action_execute(action_ctx) - mock_execute.assert_called_with("EXECUTE STREAMLIT test_streamlit_deploy_snowcli()") + mock_execute.assert_called_with(f"EXECUTE STREAMLIT {STREAMLIT_NAME}()") def test_get_execute_sql(example_streamlit_workspace): entity, action_ctx = example_streamlit_workspace - execute_sql = entity.action_get_execute_sql(action_ctx) + execute_sql = entity.get_execute_sql(action_ctx) - assert execute_sql == "EXECUTE STREAMLIT test_streamlit_deploy_snowcli()" + assert execute_sql == f"EXECUTE STREAMLIT {STREAMLIT_NAME}()" def test_get_drop_sql(example_streamlit_workspace): entity, action_ctx = example_streamlit_workspace - drop_sql = entity.action_get_drop_sql(action_ctx) + drop_sql = entity.get_drop_sql(action_ctx) + + assert drop_sql == f"DROP STREAMLIT {STREAMLIT_NAME}" + +def test_get_deploy_sql(example_streamlit_workspace): + entity, action_ctx = example_streamlit_workspace + deploy_sql = entity.get_deploy_sql(action_ctx) - assert drop_sql == "DROP STREAMLIT test_streamlit_deploy_snowcli" + assert deploy_sql == f"""CREATE STREAMLIT IDENTIFIER('{STREAMLIT_NAME}') + MAIN_FILE = 'streamlit_app.py' + QUERY_WAREHOUSE = 'test_warehouse' + TITLE = 'My Fancy Streamlit' +""" \ No newline at end of file diff --git a/tests/test_data/projects/example_streamlit_v2/snowflake.yml b/tests/test_data/projects/example_streamlit_v2/snowflake.yml index 27205e6595..f14d3e9d1b 100644 --- a/tests/test_data/projects/example_streamlit_v2/snowflake.yml +++ b/tests/test_data/projects/example_streamlit_v2/snowflake.yml @@ -2,11 +2,13 @@ definition_version: '2' entities: test_streamlit: type: "streamlit" - identifier: test_streamlit_deploy_snowcli + identifier: test_streamlit title: "My Fancy Streamlit" stage: streamlit - query_warehouse: xsmall + query_warehouse: test_warehouse main_file: streamlit_app.py - stage: streamlit artifacts: - streamlit_app.py + - utils/utils.py + - pages/ + - environment.yml