Add external access support in streamlit (#1439)

RELEASE-NOTES.md

@@ -32,6 +32,7 @@
 * Added `snow spcs service execute-job` command, which supports creating and executing a job service in the current schema.
 * Added `snow app events` command to fetch logs and traces from local and customer app installations
 * Added support for project definition file defaults in templates
+* Added support for external access (api integrations and secrets) in Streamlit.
 
 ## Fixes and improvements
 * Fixed problem with whitespaces in `snow connection add` command

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

@@ -155,17 +155,7 @@ def streamlit_deploy(
 
     # Get first streamlit
     streamlit: StreamlitEntityModel = streamlits[list(streamlits)[0]]
-    streamlit_id = streamlit.fqn.using_context()
-
-    url = StreamlitManager().deploy(
-        streamlit_id=streamlit_id,
-        artifacts=streamlit.artifacts,
-        stage_name=streamlit.stage,
-        main_file=streamlit.main_file,
-        replace=replace,
-        query_warehouse=streamlit.query_warehouse,
-        title=streamlit.title,
-    )
+    url = StreamlitManager().deploy(streamlit=streamlit, replace=replace)
 
     if open_:
         typer.launch(url)

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

@@ -29,6 +29,9 @@
 )
 from snowflake.cli.api.feature_flags import FeatureFlag
 from snowflake.cli.api.identifiers import FQN
+from snowflake.cli.api.project.schemas.entities.streamlit_entity_model import (
+    StreamlitEntityModel,
+)
 from snowflake.cli.api.sql_execution import SqlExecutionMixin
 from snowflake.connector.cursor import SnowflakeCursor
 from snowflake.connector.errors import ProgrammingError
@@ -62,14 +65,12 @@ def _put_streamlit_files(
 
     def _create_streamlit(
         self,
-        streamlit_id: FQN,
-        main_file: str,
+        streamlit: StreamlitEntityModel,
         replace: Optional[bool] = None,
         experimental: Optional[bool] = None,
-        query_warehouse: Optional[str] = None,
         from_stage_name: Optional[str] = None,
-        title: Optional[str] = None,
     ):
+        streamlit_id = streamlit.fqn.using_connection(self._conn)
         query = []
         if replace:
             query.append(f"CREATE OR REPLACE STREAMLIT {streamlit_id.sql_identifier}")
@@ -87,25 +88,24 @@ def _create_streamlit(
         if from_stage_name:
             query.append(f"ROOT_LOCATION = '{from_stage_name}'")
 
-        query.append(f"MAIN_FILE = '{main_file}'")
+        query.append(f"MAIN_FILE = '{streamlit.main_file}'")
+
+        if streamlit.query_warehouse:
+            query.append(f"QUERY_WAREHOUSE = {streamlit.query_warehouse}")
+        if streamlit.title:
+            query.append(f"TITLE = '{streamlit.title}'")
 
-        if query_warehouse:
-            query.append(f"QUERY_WAREHOUSE = {query_warehouse}")
-        if title:
-            query.append(f"TITLE = '{title}'")
+        if streamlit.external_access_integrations:
+            query.append(streamlit.get_external_access_integrations_sql())
+
+        if streamlit.secrets:
+            query.append(streamlit.get_secrets_sql())
 
         self._execute_query("\n".join(query))
 
-    def deploy(
-        self,
-        streamlit_id: FQN,
-        main_file: str,
-        artifacts: Optional[List[Path]] = None,
-        stage_name: Optional[str] = None,
-        query_warehouse: Optional[str] = None,
-        replace: Optional[bool] = False,
-        title: Optional[str] = None,
-    ):
+    def deploy(self, streamlit: StreamlitEntityModel, replace: bool = False):
+        streamlit_id = streamlit.fqn.using_connection(self._conn)
+
         # for backwards compatibility - quoted stage path might be case-sensitive
         # https://docs.snowflake.com/en/sql-reference/identifiers-syntax#double-quoted-identifiers
         streamlit_name_for_root_location = streamlit_id.name
@@ -122,12 +122,9 @@ def deploy(
             # TODO: Support from_stage
             # from_stage_stmt = f"FROM_STAGE = '{stage_name}'" if stage_name else ""
             self._create_streamlit(
-                streamlit_id,
-                main_file,
+                streamlit=streamlit,
                 replace=replace,
-                query_warehouse=query_warehouse,
                 experimental=True,
-                title=title,
             )
             try:
                 if use_versioned_stage:
@@ -157,7 +154,7 @@ def deploy(
 
             self._put_streamlit_files(
                 root_location,
-                artifacts,
+                streamlit.artifacts,
             )
         else:
             """
@@ -167,7 +164,7 @@ def deploy(
             """
             stage_manager = StageManager()
 
-            stage_name = stage_name or "streamlit"
+            stage_name = streamlit.stage or "streamlit"
             stage_name = FQN.from_string(stage_name).using_connection(self._conn)
 
             stage_manager.create(fqn=stage_name)
@@ -176,16 +173,13 @@ def deploy(
                 f"{stage_name}/{streamlit_name_for_root_location}"
             )
 
-            self._put_streamlit_files(root_location, artifacts)
+            self._put_streamlit_files(root_location, streamlit.artifacts)
 
             self._create_streamlit(
-                streamlit_id,
-                main_file,
+                streamlit=streamlit,
                 replace=replace,
-                query_warehouse=query_warehouse,
                 from_stage_name=root_location,
                 experimental=False,
-                title=title,
             )
 
         return self.get_url(streamlit_name=streamlit_id)

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

@@ -109,3 +109,31 @@ def get_type(self) -> type:
         them in __pydantic_generic_metadata__
         """
         return self.__pydantic_generic_metadata__["args"][0]
+
+
+from typing import Dict, List, Optional
+
+from pydantic import Field
+
+
+class ExternalAccessBaseModel:
+    external_access_integrations: Optional[List[str]] = Field(
+        title="Names of external access integrations needed for this entity to access external networks",
+        default=[],
+    )
+    secrets: Optional[Dict[str, str]] = Field(
+        title="Assigns the names of secrets to variables so that you can use the variables to reference the secrets",
+        default={},
+    )
+
+    def get_external_access_integrations_sql(self) -> str | None:
+        if not self.external_access_integrations:
+            return None
+        external_access_integration_name = ", ".join(self.external_access_integrations)
+        return f"external_access_integrations=({external_access_integration_name})"
+
+    def get_secrets_sql(self) -> str | None:
+        if not self.secrets:
+            return None
+        secrets = ", ".join(f"'{key}' = {value}" for key, value in self.secrets.items())
+        return f"secrets = ({secrets})"

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

@@ -14,15 +14,18 @@
 
 from __future__ import annotations
 
-from typing import Dict, List, Literal, Optional, Union
+from typing import List, Literal, Optional, Union
 
 from pydantic import Field, field_validator
-from snowflake.cli.api.project.schemas.entities.common import EntityModelBase
+from snowflake.cli.api.project.schemas.entities.common import (
+    EntityModelBase,
+    ExternalAccessBaseModel,
+)
 from snowflake.cli.api.project.schemas.snowpark.argument import Argument
 from snowflake.cli.api.project.schemas.updatable_model import DiscriminatorField
 
 
-class SnowparkEntityModel(EntityModelBase):
+class SnowparkEntityModel(EntityModelBase, ExternalAccessBaseModel):
     handler: str = Field(
         title="Function’s or procedure’s implementation of the object inside source module",
         examples=["functions.hello_function"],
@@ -36,14 +39,6 @@ class SnowparkEntityModel(EntityModelBase):
     runtime: Optional[Union[str, float]] = Field(
         title="Python version to use when executing ", default=None
     )
-    external_access_integrations: Optional[List[str]] = Field(
-        title="Names of external access integrations needed for this procedure’s handler code to access external networks",
-        default=[],
-    )
-    secrets: Optional[Dict[str, str]] = Field(
-        title="Assigns the names of secrets to variables so that you can use the variables to reference the secrets",
-        default={},
-    )
     imports: Optional[List[str]] = Field(
         title="Stage and path to previously uploaded files you want to import",
         default=[],

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

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

src/snowflake/cli/api/sql_execution.py

@@ -35,16 +35,20 @@
     unquote_identifier,
 )
 from snowflake.cli.api.utils.cursor import find_first_row
+from snowflake.connector import SnowflakeConnection
 from snowflake.connector.cursor import DictCursor, SnowflakeCursor
 from snowflake.connector.errors import ProgrammingError
 
 
 class SqlExecutionMixin:
-    def __init__(self):
+    def __init__(self, connection: SnowflakeConnection | None = None):
         self._snowpark_session = None
+        self._connection = connection
 
     @property
-    def _conn(self):
+    def _conn(self) -> SnowflakeConnection:
+        if self._connection:
+            return self._connection
         return get_cli_context().connection
 
     @property

tests/streamlit/test_streamlit_manager.py

@@ -0,0 +1,85 @@
+from pathlib import Path
+from textwrap import dedent
+from unittest import mock
+from unittest.mock import MagicMock
+
+from snowflake.cli._plugins.streamlit.manager import StreamlitManager
+from snowflake.cli.api.project.schemas.entities.streamlit_entity_model import (
+    StreamlitEntityModel,
+)
+
+
+@mock.patch("snowflake.cli._plugins.streamlit.manager.StageManager")
+@mock.patch("snowflake.cli._plugins.streamlit.manager.StreamlitManager.get_url")
+@mock.patch("snowflake.cli._plugins.streamlit.manager.StreamlitManager._execute_query")
+def test_deploy_streamlit(mock_execute_query, _, mock_stage_manager, temp_dir):
+    mock_stage_manager().get_standard_stage_prefix.return_value = "stage_root"
+
+    main_file = Path(temp_dir) / "main.py"
+    main_file.touch()
+
+    st = StreamlitEntityModel(
+        type="streamlit",
+        identifier="my_streamlit_app",
+        title="MyStreamlit",
+        query_warehouse="My_WH",
+        main_file=str(main_file),
+        # Possibly can be PathMapping
+        artifacts=[main_file],
+    )
+
+    StreamlitManager(MagicMock(database="DB", schema="SH")).deploy(
+        streamlit=st, replace=False
+    )
+
+    mock_execute_query.assert_called_once_with(
+        dedent(
+            f"""\
+        CREATE STREAMLIT IDENTIFIER('DB.SH.my_streamlit_app')
+        ROOT_LOCATION = 'stage_root'
+        MAIN_FILE = '{main_file}'
+        QUERY_WAREHOUSE = My_WH
+        TITLE = 'MyStreamlit'"""
+        )
+    )
+
+
+@mock.patch("snowflake.cli._plugins.streamlit.manager.StageManager")
+@mock.patch("snowflake.cli._plugins.streamlit.manager.StreamlitManager.get_url")
+@mock.patch("snowflake.cli._plugins.streamlit.manager.StreamlitManager._execute_query")
+def test_deploy_streamlit_with_api_integrations(
+    mock_execute_query, _, mock_stage_manager, temp_dir
+):
+    mock_stage_manager().get_standard_stage_prefix.return_value = "stage_root"
+
+    main_file = Path(temp_dir) / "main.py"
+    main_file.touch()
+
+    st = StreamlitEntityModel(
+        type="streamlit",
+        identifier="my_streamlit_app",
+        title="MyStreamlit",
+        query_warehouse="My_WH",
+        main_file=str(main_file),
+        # Possibly can be PathMapping
+        artifacts=[main_file],
+        external_access_integrations=["MY_INTERGATION", "OTHER"],
+        secrets={"my_secret": "SecretOfTheSecrets", "other": "other_secret"},
+    )
+
+    StreamlitManager(MagicMock(database="DB", schema="SH")).deploy(
+        streamlit=st, replace=False
+    )
+
+    mock_execute_query.assert_called_once_with(
+        dedent(
+            f"""\
+        CREATE STREAMLIT IDENTIFIER('DB.SH.my_streamlit_app')
+        ROOT_LOCATION = 'stage_root'
+        MAIN_FILE = '{main_file}'
+        QUERY_WAREHOUSE = My_WH
+        TITLE = 'MyStreamlit'
+        external_access_integrations=(MY_INTERGATION, OTHER)
+        secrets = ('my_secret' = SecretOfTheSecrets, 'other' = other_secret)"""
+        )
+    )

tests_integration/test_data/projects/streamlit_v2/snowflake.yml

@@ -10,6 +10,3 @@ entities:
     main_file: streamlit_app.py
     artifacts:
       - streamlit_app.py
-      - utils/utils.py
-      - pages/
-      - environment.yml

tests_integration/test_data/projects/streamlit_v2/streamlit_app.py

@@ -1,4 +1,4 @@
 import streamlit as st
-from utils.utils import hello
 
-st.title(f"Example streamlit app. {hello()}")
+
+st.title(f"Example streamlit app.")

tests_integration/test_data/projects/streamlit_v2_external_access/snowflake.yml

@@ -0,0 +1,16 @@
+definition_version: 2
+entities:
+  my_streamlit:
+    type: "streamlit"
+    identifier:
+      name: test_streamlit_deploy_snowcli_ext_access
+    title: "My Fancy Streamlit"
+    stage: streamlit
+    query_warehouse: xsmall
+    main_file: streamlit_app.py
+    external_access_integrations:
+      - snowflake_docs_access_integration
+    secrets:
+      generic_secret: external_access_db.public.test_secret
+    artifacts:
+      - streamlit_app.py

tests_integration/test_data/projects/streamlit_v2_external_access/streamlit_app.py

@@ -0,0 +1,6 @@
+import streamlit as st
+import _snowflake
+
+
+st.title(f"Example streamlit app.")
+secret = _snowflake.get_generic_secret_string("generic_secret")