Use IDENTIFIER when building streamlit queries (#1283)

RELEASE-NOTES.md

@@ -58,6 +58,7 @@
 * Added log into the file with loaded external plugins
 * Warn users if they attempt to use templating with project definition version 1
 * Improved output and format of Pydantic validation errors
+* Improved support for quoted identifiers in streamlit commands.
 * The `snow app run` command will no longer override debug mode during an application upgrade unless explicitly set in `snowflake.yml`
 
 # v2.5.0

src/snowflake/cli/api/identifiers.py

@@ -64,6 +64,10 @@ def identifier(self) -> str:
     def url_identifier(self) -> str:
         return ".".join(identifier_for_url(part) for part in self.identifier.split("."))
 
+    @property
+    def sql_identifier(self) -> str:
+        return f"IDENTIFIER('{self.identifier}')"
+
     def __str__(self):
         return self.identifier
 

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

@@ -49,8 +49,18 @@
 log = logging.getLogger(__name__)
 
 
+class IdentifierType(click.ParamType):
+    name = "TEXT"
+
+    def convert(self, value, param, ctx):
+        return FQN.from_string(value)
+
+
 StreamlitNameArgument = typer.Argument(
-    ..., help="Name of the Streamlit app.", show_default=False
+    ...,
+    help="Name of the Streamlit app.",
+    show_default=False,
+    click_type=IdentifierType(),
 )
 OpenOption = typer.Option(
     False,
@@ -79,7 +89,7 @@
 
 @app.command("share", requires_connection=True)
 def streamlit_share(
-    name: str = StreamlitNameArgument,
+    name: FQN = StreamlitNameArgument,
     to_role: str = typer.Argument(
         ..., help="Role with which to share the Streamlit app."
     ),
@@ -138,10 +148,10 @@ def streamlit_deploy(
     elif pages_dir is None:
         pages_dir = "pages"
 
-    streamlit_name = FQN.from_identifier_model(streamlit).using_context()
+    streamlit_id = FQN.from_identifier_model(streamlit).using_context()
 
     url = StreamlitManager().deploy(
-        streamlit=streamlit_name,
+        streamlit_id=streamlit_id,
         environment_file=Path(environment_file),
         pages_dir=Path(pages_dir),
         stage_name=streamlit.stage,
@@ -161,7 +171,7 @@ def streamlit_deploy(
 
 @app.command("get-url", requires_connection=True)
 def get_url(
-    name: str = StreamlitNameArgument,
+    name: FQN = StreamlitNameArgument,
     open_: bool = OpenOption,
     **options,
 ):

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

@@ -36,9 +36,9 @@
 
 
 class StreamlitManager(SqlExecutionMixin):
-    def share(self, streamlit_name: str, to_role: str) -> SnowflakeCursor:
+    def share(self, streamlit_name: FQN, to_role: str) -> SnowflakeCursor:
         return self._execute_query(
-            f"grant usage on streamlit {streamlit_name} to role {to_role}"
+            f"grant usage on streamlit {streamlit_name.sql_identifier} to role {to_role}"
         )
 
     def _put_streamlit_files(
@@ -72,7 +72,7 @@ def _put_streamlit_files(
 
     def _create_streamlit(
         self,
-        fully_qualified_name: str,
+        streamlit_id: FQN,
         main_file: Path,
         replace: Optional[bool] = None,
         experimental: Optional[bool] = None,
@@ -82,15 +82,17 @@ def _create_streamlit(
     ):
         query = []
         if replace:
-            query.append(f"CREATE OR REPLACE STREAMLIT {fully_qualified_name}")
+            query.append(f"CREATE OR REPLACE STREAMLIT {streamlit_id.sql_identifier}")
         elif experimental:
             # For experimental behaviour, we need to use CREATE STREAMLIT IF NOT EXISTS
             # for a streamlit app with an embedded stage
             # because this is analogous to the behavior for non-experimental
             # deploy which does CREATE STAGE IF NOT EXISTS
-            query.append(f"CREATE STREAMLIT IF NOT EXISTS {fully_qualified_name}")
+            query.append(
+                f"CREATE STREAMLIT IF NOT EXISTS {streamlit_id.sql_identifier}"
+            )
         else:
-            query.append(f"CREATE STREAMLIT {fully_qualified_name}")
+            query.append(f"CREATE STREAMLIT {streamlit_id.sql_identifier}")
 
         if from_stage_name:
             query.append(f"ROOT_LOCATION = '{from_stage_name}'")
@@ -106,7 +108,7 @@ def _create_streamlit(
 
     def deploy(
         self,
-        streamlit: FQN,
+        streamlit_id: FQN,
         main_file: Path,
         environment_file: Optional[Path] = None,
         pages_dir: Optional[Path] = None,
@@ -119,8 +121,7 @@ def deploy(
     ):
         # 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.name
-        fully_qualified_name = streamlit.identifier
+        streamlit_name_for_root_location = streamlit_id.name
 
         if (
             experimental_behaviour_enabled()
@@ -133,23 +134,25 @@ def deploy(
             # TODO: Support from_stage
             # from_stage_stmt = f"FROM_STAGE = '{stage_name}'" if stage_name else ""
             self._create_streamlit(
-                fully_qualified_name,
+                streamlit_id,
                 main_file,
                 replace=replace,
                 query_warehouse=query_warehouse,
                 experimental=True,
                 title=title,
             )
             try:
-                self._execute_query(f"ALTER streamlit {fully_qualified_name} CHECKOUT")
+                self._execute_query(
+                    f"ALTER streamlit {streamlit_id.identifier} CHECKOUT"
+                )
             except ProgrammingError as e:
                 # If an error is raised because a CHECKOUT has already occurred,
                 # simply skip it and continue
                 if "Checkout already exists" in str(e):
                     log.info("Checkout already exists, continuing")
                 else:
                     raise
-            stage_path = streamlit.identifier
+            stage_path = streamlit_id.identifier
             embedded_stage_name = f"snow://streamlit/{stage_path}"
             root_location = f"{embedded_stage_name}/default_checkout"
 
@@ -186,7 +189,7 @@ def deploy(
             )
 
             self._create_streamlit(
-                fully_qualified_name,
+                streamlit_id,
                 main_file,
                 replace=replace,
                 query_warehouse=query_warehouse,
@@ -195,11 +198,11 @@ def deploy(
                 title=title,
             )
 
-        return self.get_url(fully_qualified_name)
+        return self.get_url(streamlit_name=streamlit_id)
 
-    def get_url(self, streamlit_name: str, database=None, schema=None) -> str:
+    def get_url(self, streamlit_name: FQN) -> str:
         try:
-            fqn = FQN.from_string(streamlit_name).using_connection(self._conn)
+            fqn = streamlit_name.using_connection(self._conn)
             return make_snowsight_url(
                 self._conn,
                 f"/#/streamlit-apps/{fqn.url_identifier}",

tests/streamlit/test_commands.py

@@ -91,7 +91,7 @@ def test_deploy_only_streamlit_file(
         ),
         dedent(
             f"""
-            CREATE STREAMLIT MockDatabase.MockSchema.{STREAMLIT_NAME}
+            CREATE STREAMLIT IDENTIFIER('MockDatabase.MockSchema.{STREAMLIT_NAME}')
             ROOT_LOCATION = '@MockDatabase.MockSchema.streamlit/{STREAMLIT_NAME}'
             MAIN_FILE = 'streamlit_app.py'
             QUERY_WAREHOUSE = test_warehouse
@@ -142,7 +142,7 @@ def test_deploy_only_streamlit_file_no_stage(
         ),
         dedent(
             f"""
-            CREATE STREAMLIT MockDatabase.MockSchema.{STREAMLIT_NAME}
+            CREATE STREAMLIT IDENTIFIER('MockDatabase.MockSchema.{STREAMLIT_NAME}')
             ROOT_LOCATION = '@MockDatabase.MockSchema.streamlit/{STREAMLIT_NAME}'
             MAIN_FILE = 'streamlit_app.py'
             QUERY_WAREHOUSE = test_warehouse
@@ -192,7 +192,7 @@ def test_deploy_only_streamlit_file_replace(
         ),
         dedent(
             f"""
-            CREATE OR REPLACE STREAMLIT MockDatabase.MockSchema.{STREAMLIT_NAME}
+            CREATE OR REPLACE STREAMLIT IDENTIFIER('MockDatabase.MockSchema.{STREAMLIT_NAME}')
             ROOT_LOCATION = '@MockDatabase.MockSchema.streamlit/{STREAMLIT_NAME}'
             MAIN_FILE = 'streamlit_app.py'
             QUERY_WAREHOUSE = test_warehouse
@@ -260,7 +260,7 @@ def test_deploy_streamlit_and_environment_files(
         _put_query("environment.yml", root_path),
         dedent(
             f"""
-            CREATE STREAMLIT MockDatabase.MockSchema.{STREAMLIT_NAME}
+            CREATE STREAMLIT IDENTIFIER('MockDatabase.MockSchema.{STREAMLIT_NAME}')
             ROOT_LOCATION = '@MockDatabase.MockSchema.streamlit/{STREAMLIT_NAME}'
             MAIN_FILE = 'streamlit_app.py'
             QUERY_WAREHOUSE = test_warehouse
@@ -301,7 +301,7 @@ def test_deploy_streamlit_and_pages_files(
         _put_query("pages/*.py", f"{root_path}/pages"),
         dedent(
             f"""
-            CREATE STREAMLIT MockDatabase.MockSchema.{STREAMLIT_NAME}
+            CREATE STREAMLIT IDENTIFIER('MockDatabase.MockSchema.{STREAMLIT_NAME}')
             ROOT_LOCATION = '@MockDatabase.MockSchema.streamlit/{STREAMLIT_NAME}'
             MAIN_FILE = 'streamlit_app.py'
             QUERY_WAREHOUSE = test_warehouse
@@ -344,7 +344,7 @@ def test_deploy_all_streamlit_files(
         _put_query("extra_file.py", root_path),
         dedent(
             f"""
-            CREATE STREAMLIT MockDatabase.MockSchema.{STREAMLIT_NAME}
+            CREATE STREAMLIT IDENTIFIER('MockDatabase.MockSchema.{STREAMLIT_NAME}')
             ROOT_LOCATION = '@MockDatabase.MockSchema.streamlit/{STREAMLIT_NAME}'
             MAIN_FILE = 'streamlit_app.py'
             QUERY_WAREHOUSE = test_warehouse
@@ -387,7 +387,7 @@ def test_deploy_put_files_on_stage(
         _put_query("pages/*.py", f"{root_path}/pages"),
         dedent(
             f"""
-            CREATE STREAMLIT MockDatabase.MockSchema.{STREAMLIT_NAME}
+            CREATE STREAMLIT IDENTIFIER('MockDatabase.MockSchema.{STREAMLIT_NAME}')
             ROOT_LOCATION = '@MockDatabase.MockSchema.streamlit_stage/{STREAMLIT_NAME}'
             MAIN_FILE = 'streamlit_app.py'
             QUERY_WAREHOUSE = test_warehouse
@@ -428,7 +428,7 @@ def test_deploy_all_streamlit_files_not_defaults(
         _put_query("streamlit_pages/*.py", f"{root_path}/pages"),
         dedent(
             f"""
-            CREATE STREAMLIT MockDatabase.MockSchema.{STREAMLIT_NAME}
+            CREATE STREAMLIT IDENTIFIER('MockDatabase.MockSchema.{STREAMLIT_NAME}')
             ROOT_LOCATION = '@MockDatabase.MockSchema.streamlit_stage/{STREAMLIT_NAME}'
             MAIN_FILE = 'main.py'
             QUERY_WAREHOUSE = streamlit_warehouse
@@ -466,7 +466,7 @@ def test_deploy_streamlit_main_and_pages_files_experimental(
     assert ctx.get_queries() == [
         dedent(
             f"""
-            CREATE STREAMLIT IF NOT EXISTS MockDatabase.MockSchema.{STREAMLIT_NAME}
+            CREATE STREAMLIT IF NOT EXISTS IDENTIFIER('MockDatabase.MockSchema.{STREAMLIT_NAME}')
             MAIN_FILE = 'streamlit_app.py'
             QUERY_WAREHOUSE = test_warehouse
             TITLE = 'My Fancy Streamlit'
@@ -532,7 +532,7 @@ def test_deploy_streamlit_main_and_pages_files_experimental_double_deploy(
     assert ctx.get_queries() == [
         dedent(
             f"""
-            CREATE STREAMLIT IF NOT EXISTS MockDatabase.MockSchema.{STREAMLIT_NAME}
+            CREATE STREAMLIT IF NOT EXISTS IDENTIFIER('MockDatabase.MockSchema.{STREAMLIT_NAME}')
             MAIN_FILE = 'streamlit_app.py'
             QUERY_WAREHOUSE = test_warehouse
             TITLE = 'My Fancy Streamlit'
@@ -573,7 +573,7 @@ def test_deploy_streamlit_main_and_pages_files_experimental_no_stage(
     assert ctx.get_queries() == [
         dedent(
             f"""
-            CREATE STREAMLIT IF NOT EXISTS MockDatabase.MockSchema.{STREAMLIT_NAME}
+            CREATE STREAMLIT IF NOT EXISTS IDENTIFIER('MockDatabase.MockSchema.{STREAMLIT_NAME}')
             MAIN_FILE = 'streamlit_app.py'
             QUERY_WAREHOUSE = test_warehouse
             """
@@ -614,7 +614,7 @@ def test_deploy_streamlit_main_and_pages_files_experimental_replace(
     assert ctx.get_queries() == [
         dedent(
             f"""
-            CREATE OR REPLACE STREAMLIT MockDatabase.MockSchema.{STREAMLIT_NAME}
+            CREATE OR REPLACE STREAMLIT IDENTIFIER('MockDatabase.MockSchema.{STREAMLIT_NAME}')
             MAIN_FILE = 'streamlit_app.py'
             QUERY_WAREHOUSE = test_warehouse
             TITLE = 'My Fancy Streamlit'
@@ -662,7 +662,8 @@ def test_share_streamlit(mock_connector, runner, mock_ctx):
 
     assert result.exit_code == 0, result.output
     assert (
-        ctx.get_query() == f"grant usage on streamlit {STREAMLIT_NAME} to role {role}"
+        ctx.get_query()
+        == f"grant usage on streamlit IDENTIFIER('{STREAMLIT_NAME}') to role {role}"
     )