Add streamlit entities (#1934)

* Actions

* fix

* text fix

* Fixes

* Fix
Co-authored-by: Patryk Czajka <[email protected]>

src/snowflake/cli/_plugins/nativeapp/entities/application_package.py

@@ -966,6 +966,8 @@ def _bundle_children(self, action_ctx: ActionContext) -> List[str]:
                 child_entity.get_deploy_sql(
                     artifacts_dir=child_artifacts_dir.relative_to(self.deploy_root),
                     schema=child_schema,
+                    # TODO Allow users to override the hard-coded value for specific children
+                    replace=True,
                 )
             )
             if app_role:

src/snowflake/cli/_plugins/nativeapp/feature_flags.py

@@ -18,7 +18,9 @@
 
 
 @unique
-class FeatureFlag(FeatureFlagMixin):
+class FeatureFlag(
+    FeatureFlagMixin
+):  # TODO move this to snowflake.cli.api.feature_flags
     ENABLE_NATIVE_APP_PYTHON_SETUP = BooleanFlag(
         "ENABLE_NATIVE_APP_PYTHON_SETUP", False
     )

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

@@ -1,23 +1,20 @@
+import functools
 from pathlib import Path
 from typing import Optional
 
-from snowflake.cli._plugins.nativeapp.artifacts import build_bundle
-from snowflake.cli._plugins.nativeapp.entities.application_package_child_interface import (
-    ApplicationPackageChildInterface,
-)
+from click import ClickException
+from snowflake.cli._plugins.connection.util import make_snowsight_url
 from snowflake.cli._plugins.nativeapp.feature_flags import FeatureFlag
 from snowflake.cli._plugins.streamlit.streamlit_entity_model import (
     StreamlitEntityModel,
 )
-from snowflake.cli.api.entities.common import EntityBase
-from snowflake.cli.api.project.schemas.v1.native_app.path_mapping import PathMapping
+from snowflake.cli._plugins.workspace.context import ActionContext
+from snowflake.cli.api.entities.common import EntityBase, get_sql_executor
+from snowflake.cli.api.secure_path import SecurePath
+from snowflake.connector.cursor import SnowflakeCursor
 
 
-# WARNING: This entity is not implemented yet. The logic below is only for demonstrating the
-# required interfaces for composability (used by ApplicationPackageEntity behind a feature flag).
-class StreamlitEntity(
-    EntityBase[StreamlitEntityModel], ApplicationPackageChildInterface
-):
+class StreamlitEntity(EntityBase[StreamlitEntityModel]):
     """
     A Streamlit app.
     """
@@ -28,43 +25,140 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
     @property
-    def project_root(self) -> Path:
+    def root(self):
         return self._workspace_ctx.project_root
 
     @property
-    def deploy_root(self) -> Path:
-        return self.project_root / "output" / "deploy"
+    def artifacts(self):
+        return self._entity_model.artifacts
 
-    def action_bundle(
-        self,
-        *args,
-        **kwargs,
-    ):
+    @functools.cached_property
+    def _sql_executor(self):
+        return get_sql_executor()
+
+    @functools.cached_property
+    def _conn(self):
+        return self._sql_executor._conn  # noqa
+
+    @property
+    def model(self):
+        return self._entity_model  # noqa
+
+    def action_bundle(self, action_ctx: ActionContext, *args, **kwargs):
         return self.bundle()
 
-    def bundle(self, bundle_root=None):
-        return build_bundle(
-            self.project_root,
-            bundle_root or self.deploy_root,
-            [
-                PathMapping(src=str(artifact))
-                for artifact in self._entity_model.artifacts
-            ],
+    def action_deploy(self, action_ctx: ActionContext, *args, **kwargs):
+        # After adding bundle map- we should use it's mapping here
+        # To copy artifacts to destination on stage.
+
+        return self._sql_executor.execute_query(self.get_deploy_sql())
+
+    def action_drop(self, action_ctx: ActionContext, *args, **kwargs):
+        return self._sql_executor.execute_query(self.get_drop_sql())
+
+    def action_execute(
+        self, action_ctx: ActionContext, *args, **kwargs
+    ) -> SnowflakeCursor:
+        return self._sql_executor.execute_query(self.get_execute_sql())
+
+    def action_get_url(
+        self, action_ctx: ActionContext, *args, **kwargs
+    ):  # maybe this should be a property
+        name = self._entity_model.fqn.using_connection(self._conn)
+        return make_snowsight_url(
+            self._conn, f"/#/streamlit-apps/{name.url_identifier}"
         )
 
+    def bundle(self, output_dir: Optional[Path] = None):
+
+        if not output_dir:
+            output_dir = self.root / "output" / self._entity_model.stage
+
+        artifacts = self._entity_model.artifacts
+
+        output_dir.mkdir(parents=True, exist_ok=True)  # type: ignore
+
+        output_files = []
+
+        # This is far from , but will be replaced by bundlemap mappings.
+        for file in artifacts:
+            output_file = output_dir / file.name
+
+            if file.is_file():
+                SecurePath(file).copy(output_file)
+            elif file.is_dir():
+                output_file.mkdir(parents=True, exist_ok=True)
+                SecurePath(file).copy(output_file, dirs_exist_ok=True)
+
+                output_files.append(output_file)
+
+        return output_files
+
+    def action_share(
+        self, action_ctx: ActionContext, to_role: str, *args, **kwargs
+    ) -> SnowflakeCursor:
+        return self._sql_executor.execute_query(self.get_share_sql(to_role))
+
     def get_deploy_sql(
         self,
+        if_not_exists: bool = False,
+        replace: bool = False,
+        from_stage_name: Optional[str] = None,
         artifacts_dir: Optional[Path] = None,
         schema: Optional[str] = None,
+        *args,
+        **kwargs,
     ):
-        entity_id = self.entity_id
-        if artifacts_dir:
-            streamlit_name = f"{schema}.{entity_id}" if schema else entity_id
-            return f"CREATE OR REPLACE STREAMLIT {streamlit_name} FROM '{artifacts_dir}' MAIN_FILE='{self._entity_model.main_file}';"
+        if replace and if_not_exists:
+            raise ClickException("Cannot specify both replace and if_not_exists")
+
+        if replace:
+            query = "CREATE OR REPLACE "
+        elif if_not_exists:
+            query = "CREATE IF NOT EXISTS "
         else:
-            return f"CREATE OR REPLACE STREAMLIT {entity_id} MAIN_FILE='{self._entity_model.main_file}';"
+            query = "CREATE "
+
+        schema_to_use = schema or self._entity_model.fqn.schema
+        query += f"STREAMLIT {self._entity_model.fqn.set_schema(schema_to_use).sql_identifier}"
+
+        if from_stage_name:
+            query += f"\nROOT_LOCATION = '{from_stage_name}'"
+        elif artifacts_dir:
+            query += f"\nFROM '{artifacts_dir}'"
+
+        query += f"\nMAIN_FILE = '{self._entity_model.main_file}'"
+
+        if self.model.imports:
+            query += "\n" + self.model.get_imports_sql()
+
+        if self.model.query_warehouse:
+            query += f"\nQUERY_WAREHOUSE = '{self.model.query_warehouse}'"
+
+        if self.model.title:
+            query += f"\nTITLE = '{self.model.title}'"
+
+        if self.model.comment:
+            query += f"\nCOMMENT = '{self.model.comment}'"
+
+        if self.model.external_access_integrations:
+            query += "\n" + self.model.get_external_access_integrations_sql()
+
+        if self.model.secrets:
+            query += "\n" + self.model.get_secrets_sql()
+
+        return query + ";"
+
+    def get_drop_sql(self):
+        return f"DROP STREAMLIT {self._entity_model.fqn};"
+
+    def get_execute_sql(self):
+        return f"EXECUTE STREAMLIT {self._entity_model.fqn}();"
+
+    def get_share_sql(self, to_role: str) -> str:
+        return f"GRANT USAGE ON STREAMLIT {self.model.fqn.sql_identifier} TO ROLE {to_role};"
 
-    def get_usage_grant_sql(self, app_role: str, schema: Optional[str] = None):
+    def get_usage_grant_sql(self, app_role: str, schema: Optional[str] = None) -> str:
         entity_id = self.entity_id
         streamlit_name = f"{schema}.{entity_id}" if schema else entity_id
         return (

tests/nativeapp/test_children.py

@@ -143,10 +143,12 @@ def test_children_bundle_with_custom_dir(project_directory):
                 dedent(
                     f"""
                     -- AUTO GENERATED CHILDREN SECTION
-                    CREATE OR REPLACE STREAMLIT v_schema.my_streamlit FROM '{custom_dir_path}' MAIN_FILE='streamlit_app.py';
+                    CREATE OR REPLACE STREAMLIT IDENTIFIER('v_schema.my_streamlit')
+                    FROM '{custom_dir_path}'
+                    MAIN_FILE = 'streamlit_app.py';
                     CREATE APPLICATION ROLE IF NOT EXISTS my_app_role;
                     GRANT USAGE ON SCHEMA v_schema TO APPLICATION ROLE my_app_role;
                     GRANT USAGE ON STREAMLIT v_schema.my_streamlit TO APPLICATION ROLE my_app_role;
-                    """
+"""
                 )
             )

tests/streamlit/__snapshots__/test_actions.ambr

@@ -0,0 +1,49 @@
+# serializer version: 1
+# name: test_get_deploy_sql[kwargs0]
+  '''
+  CREATE OR REPLACE STREAMLIT IDENTIFIER('test_streamlit') 
+   MAIN_FILE = 'streamlit_app.py' 
+   QUERY_WAREHOUSE = 'test_warehouse' 
+   TITLE = 'My Fancy Streamlit' 
+
+  '''
+# ---
+# name: test_get_deploy_sql[kwargs1]
+  '''
+  CREATE IF NOT EXISTS STREAMLIT IDENTIFIER('test_streamlit') 
+   MAIN_FILE = 'streamlit_app.py' 
+   QUERY_WAREHOUSE = 'test_warehouse' 
+   TITLE = 'My Fancy Streamlit' 
+
+  '''
+# ---
+# name: test_get_deploy_sql[kwargs2]
+  '''
+  CREATE STREAMLIT IDENTIFIER('test_streamlit') 
+   ROOT_LOCATION = 'test_stage' 
+   MAIN_FILE = 'streamlit_app.py' 
+   QUERY_WAREHOUSE = 'test_warehouse' 
+   TITLE = 'My Fancy Streamlit' 
+
+  '''
+# ---
+# name: test_get_deploy_sql[kwargs3]
+  '''
+  CREATE OR REPLACE STREAMLIT IDENTIFIER('test_streamlit') 
+   ROOT_LOCATION = 'test_stage' 
+   MAIN_FILE = 'streamlit_app.py' 
+   QUERY_WAREHOUSE = 'test_warehouse' 
+   TITLE = 'My Fancy Streamlit' 
+
+  '''
+# ---
+# name: test_get_deploy_sql[kwargs4]
+  '''
+  CREATE IF NOT EXISTS STREAMLIT IDENTIFIER('test_streamlit') 
+   ROOT_LOCATION = 'test_stage' 
+   MAIN_FILE = 'streamlit_app.py' 
+   QUERY_WAREHOUSE = 'test_warehouse' 
+   TITLE = 'My Fancy Streamlit' 
+
+  '''
+# ---

tests/streamlit/__snapshots__/test_commands.ambr

@@ -1,5 +1,5 @@
 # serializer version: 1
-# name: test_artifacts_must_exists
+# name: test_artifacts_must_exist
   '''
   +- Error ----------------------------------------------------------------------+
   | During evaluation of DefinitionV20 in project definition following errors    |
@@ -11,38 +11,6 @@
 
   '''
 # ---
-# 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 ----------------------------------------------------------------------+
@@ -74,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.      |

tests/streamlit/__snapshots__/test_streamlit_entity.ambr

@@ -0,0 +1,60 @@
+# serializer version: 1
+# name: test_get_deploy_sql[kwargs0]
+  '''
+  CREATE OR REPLACE STREAMLIT IDENTIFIER('test_streamlit')
+  MAIN_FILE = 'streamlit_app.py'
+  QUERY_WAREHOUSE = 'test_warehouse'
+  TITLE = 'My Fancy Streamlit';
+  '''
+# ---
+# name: test_get_deploy_sql[kwargs1]
+  '''
+  CREATE IF NOT EXISTS STREAMLIT IDENTIFIER('test_streamlit')
+  MAIN_FILE = 'streamlit_app.py'
+  QUERY_WAREHOUSE = 'test_warehouse'
+  TITLE = 'My Fancy Streamlit';
+  '''
+# ---
+# name: test_get_deploy_sql[kwargs2]
+  '''
+  CREATE STREAMLIT IDENTIFIER('test_streamlit')
+  ROOT_LOCATION = 'test_stage'
+  MAIN_FILE = 'streamlit_app.py'
+  QUERY_WAREHOUSE = 'test_warehouse'
+  TITLE = 'My Fancy Streamlit';
+  '''
+# ---
+# name: test_get_deploy_sql[kwargs3]
+  '''
+  CREATE OR REPLACE STREAMLIT IDENTIFIER('test_streamlit')
+  ROOT_LOCATION = 'test_stage'
+  MAIN_FILE = 'streamlit_app.py'
+  QUERY_WAREHOUSE = 'test_warehouse'
+  TITLE = 'My Fancy Streamlit';
+  '''
+# ---
+# name: test_get_deploy_sql[kwargs4]
+  '''
+  CREATE IF NOT EXISTS STREAMLIT IDENTIFIER('test_streamlit')
+  ROOT_LOCATION = 'test_stage'
+  MAIN_FILE = 'streamlit_app.py'
+  QUERY_WAREHOUSE = 'test_warehouse'
+  TITLE = 'My Fancy Streamlit';
+  '''
+# ---
+# name: test_nativeapp_children_interface
+  '''
+  CREATE STREAMLIT IDENTIFIER('test_streamlit')
+  MAIN_FILE = 'streamlit_app.py'
+  QUERY_WAREHOUSE = 'test_warehouse'
+  TITLE = 'My Fancy Streamlit';
+  '''
+# ---
+# name: test_nativeapp_children_interface.1
+  '''
+  CREATE STREAMLIT IDENTIFIER('test_streamlit')
+  MAIN_FILE = 'streamlit_app.py'
+  QUERY_WAREHOUSE = 'test_warehouse'
+  TITLE = 'My Fancy Streamlit';
+  '''
+# ---