Add imports support in Streamlit (#1590)

RELEASE-NOTES.md

@@ -44,6 +44,7 @@
 * Added templates expansion of arbitrary files for Native Apps through `templates` processor.
 * Added `SNOWFLAKE_..._PRIVATE_KEY_RAW` environment variable to pass private key as a raw string.
 * Added periodic check for newest version of Snowflake CLI. When new version is available, user will be notified.
+* Added support for `imports` in Streamlit definition.
 
 ## Fixes and improvements
 * Fixed problem with whitespaces in `snow connection add` command.

src/snowflake/cli/_plugins/snowpark/snowpark_entity_model.py

@@ -22,6 +22,7 @@
 from snowflake.cli.api.project.schemas.entities.common import (
  EntityModelBase,
  ExternalAccessBaseModel,
+ ImportsBaseModel,
 )
 from snowflake.cli.api.project.schemas.updatable_model import (
  DiscriminatorField,
@@ -43,7 +44,7 @@ class Config:
  )
 
 
-class SnowparkEntityModel(EntityModelBase, ExternalAccessBaseModel):
+class SnowparkEntityModel(EntityModelBase, ExternalAccessBaseModel, ImportsBaseModel):
  handler: str = Field(
  title="Function’s or procedure’s implementation of the object inside source module",
  examples=["functions.hello_function"],
@@ -57,10 +58,6 @@ class SnowparkEntityModel(EntityModelBase, ExternalAccessBaseModel):
  runtime: Optional[Union[str, float]] = Field(
  title="Python version to use when executing ", default=None
  )
- imports: Optional[List[str]] = Field(
- title="Stage and path to previously uploaded files you want to import",
- default=[],
- )
  stage: str = Field(title="Stage in which artifacts will be stored")
  artifacts: List[Union[PathMapping, str]] = Field(title="List of required sources")
 

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

@@ -98,7 +98,8 @@ def _create_streamlit(
  query.append(f"ROOT_LOCATION = '{from_stage_name}'")
 
  query.append(f"MAIN_FILE = '{streamlit.main_file}'")
-
+ if streamlit.imports:
+ query.append(streamlit.get_imports_sql())
  if streamlit.query_warehouse:
  query.append(f"QUERY_WAREHOUSE = {streamlit.query_warehouse}")
  if streamlit.title:

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

@@ -20,13 +20,14 @@
 from snowflake.cli.api.project.schemas.entities.common import (
  EntityModelBase,
  ExternalAccessBaseModel,
+ ImportsBaseModel,
 )
 from snowflake.cli.api.project.schemas.updatable_model import (
  DiscriminatorField,
 )
 
 
-class StreamlitEntityModel(EntityModelBase, ExternalAccessBaseModel):
+class StreamlitEntityModel(EntityModelBase, ExternalAccessBaseModel, ImportsBaseModel):
  type: Literal["streamlit"] = DiscriminatorField() # noqa: A003
  title: Optional[str] = Field(
  title="Human-readable title for the Streamlit dashboard", default=None

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

@@ -117,6 +117,19 @@ def get_type(self) -> type:
  return self.__pydantic_generic_metadata__["args"][0]
 
 
+class ImportsBaseModel:
+ imports: Optional[List[str]] = Field(
+ title="Stage and path to previously uploaded files you want to import",
+ default=[],
+ )
+
+ def get_imports_sql(self) -> str | None:
+ if not self.imports:
+ return None
+ imports = ", ".join(f"'{i}'" for i in self.imports)
+ return f"IMPORTS = ({imports})"
+
+
 class ExternalAccessBaseModel:
  external_access_integrations: Optional[List[str]] = Field(
  title="Names of external access integrations needed for this entity to access external networks",

tests/streamlit/test_streamlit_manager.py

@@ -30,6 +30,7 @@ def test_deploy_streamlit(mock_execute_query, _, mock_stage_manager, temp_dir):
  title="MyStreamlit",
  query_warehouse="My_WH",
  main_file=str(main_file),
+ imports=["@stage/foo.py", "@stage/bar.py"],
  # Possibly can be PathMapping
  artifacts=[main_file],
  )
@@ -44,6 +45,7 @@ def test_deploy_streamlit(mock_execute_query, _, mock_stage_manager, temp_dir):
  CREATE STREAMLIT IDENTIFIER('DB.SH.my_streamlit_app')
  ROOT_LOCATION = 'stage_root'
  MAIN_FILE = '{main_file}'
+ IMPORTS = ('@stage/foo.py', '@stage/bar.py')
  QUERY_WAREHOUSE = My_WH
  TITLE = 'MyStreamlit'"""
  )

tests_integration/test_streamlit.py

@@ -106,6 +106,35 @@ def test_streamlit_deploy(
  assert row_from_snowflake_session(expect) == []
 
 
+@pytest.mark.integration
+def test_streamlit_deploy_with_imports(
+ runner,
+ snowflake_session,
+ test_database,
+ _new_streamlit_role,
+ project_directory,
+ alter_snowflake_yml,
+):
+
+ # This work because uploading the imports artifact because
+ # deploying streamlit does not start the app.
+ with project_directory(f"streamlit_v2"):
+ alter_snowflake_yml(
+ "snowflake.yml",
+ "entities.my_streamlit.imports",
+ ["@stage/foo.py", "@stage/bar.py"],
+ )
+ result = runner.invoke_with_connection_json(
+ ["streamlit", "deploy", "--replace"]
+ )
+ assert result.exit_code == 0
+
+ result = runner.invoke_with_connection_json(
+ ["streamlit", "describe", "test_streamlit_deploy_snowcli"]
+ )
+ assert result.json[0]["import_urls"] == '["@stage/foo.py","@stage/bar.py"]'
+
+
 @pytest.mark.integration
 @pytest.mark.skip(
  reason="only works in accounts with experimental checkout behavior enabled"