snowflake-cli-labs 2.8.0rc1__py3-none-any.whl → 3.0.0__py3-none-any.whl

snowflake/cli/{plugins → _plugins}/streamlit/manager.py

@@ -15,22 +15,27 @@
 from __future__ import annotations
 
 import logging
-import os
 from pathlib import Path
 from typing import List, Optional
 
+from click import ClickException
+from snowflake.cli._plugins.connection.util import (
+    MissingConnectionAccountError,
+    MissingConnectionRegionError,
+    make_snowsight_url,
+)
+from snowflake.cli._plugins.object.manager import ObjectManager
+from snowflake.cli._plugins.stage.manager import StageManager
+from snowflake.cli._plugins.streamlit.streamlit_entity_model import (
+    StreamlitEntityModel,
+)
 from snowflake.cli.api.commands.experimental_behaviour import (
     experimental_behaviour_enabled,
 )
+from snowflake.cli.api.console import cli_console
 from snowflake.cli.api.feature_flags import FeatureFlag
 from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.sql_execution import SqlExecutionMixin
-from snowflake.cli.plugins.connection.util import (
-    MissingConnectionAccountError,
-    MissingConnectionRegionError,
-    make_snowsight_url,
-)
-from snowflake.cli.plugins.stage.manager import StageManager
 from snowflake.connector.cursor import SnowflakeCursor
 from snowflake.connector.errors import ProgrammingError
 
@@ -38,6 +43,10 @@ log = logging.getLogger(__name__)
 
 
 class StreamlitManager(SqlExecutionMixin):
+    def execute(self, app_name: FQN):
+        query = f"EXECUTE STREAMLIT {app_name.sql_identifier}()"
+        return self._execute_query(query=query)
+
     def share(self, streamlit_name: FQN, to_role: str) -> SnowflakeCursor:
         return self._execute_query(
             f"grant usage on streamlit {streamlit_name.sql_identifier} to role {to_role}"
@@ -46,39 +55,35 @@ class StreamlitManager(SqlExecutionMixin):
     def _put_streamlit_files(
         self,
         root_location: str,
-        main_file: Path,
-        environment_file: Optional[Path],
-        pages_dir: Optional[Path],
-        additional_source_files: Optional[List[Path]],
+        artifacts: Optional[List[Path]] = None,
     ):
+        cli_console.step(f"Deploying files to {root_location}")
+        if not artifacts:
+            return
         stage_manager = StageManager()
-
-        stage_manager.put(main_file, root_location, 4, True)
-
-        if environment_file and environment_file.exists():
-            stage_manager.put(environment_file, root_location, 4, True)
-
-        if pages_dir and pages_dir.exists():
-            stage_manager.put(pages_dir / "*.py", f"{root_location}/pages", 4, True)
-
-        if additional_source_files:
-            for file in additional_source_files:
-                if os.sep in str(file):
-                    destination = f"{root_location}/{str(file.parent)}"
-                else:
-                    destination = root_location
-                stage_manager.put(file, destination, 4, True)
+        for file in artifacts:
+            if file.is_dir():
+                if not any(file.iterdir()):
+                    cli_console.warning(f"Skipping empty directory: {file}")
+                    continue
+
+                stage_manager.put(
+                    f"{file.joinpath('*')}", f"{root_location}/{file}", 4, True
+                )
+            elif len(file.parts) > 1:
+                stage_manager.put(file, f"{root_location}/{file.parent}", 4, True)
+            else:
+                stage_manager.put(file, root_location, 4, True)
 
     def _create_streamlit(
         self,
-        streamlit_id: FQN,
-        main_file: Path,
+        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)
+        cli_console.step(f"Creating {streamlit_id} Streamlit")
         query = []
         if replace:
             query.append(f"CREATE OR REPLACE STREAMLIT {streamlit_id.sql_identifier}")
@@ -96,28 +101,35 @@ class StreamlitManager(SqlExecutionMixin):
         if from_stage_name:
             query.append(f"ROOT_LOCATION = '{from_stage_name}'")
 
-        query.append(f"MAIN_FILE = '{main_file.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:
+            query.append(f"TITLE = '{streamlit.title}'")
+
+        if streamlit.comment:
+            query.append(f"COMMENT = '{streamlit.comment}'")
+
+        if streamlit.external_access_integrations:
+            query.append(streamlit.get_external_access_integrations_sql())
 
-        if query_warehouse:
-            query.append(f"QUERY_WAREHOUSE = {query_warehouse}")
-        if title:
-            query.append(f"TITLE = '{title}'")
+        if streamlit.secrets:
+            query.append(streamlit.get_secrets_sql())
 
         self._execute_query("\n".join(query))
 
-    def deploy(
-        self,
-        streamlit_id: FQN,
-        main_file: Path,
-        environment_file: Optional[Path] = None,
-        pages_dir: Optional[Path] = None,
-        stage_name: Optional[str] = None,
-        query_warehouse: Optional[str] = None,
-        replace: Optional[bool] = False,
-        additional_source_files: Optional[List[Path]] = None,
-        title: Optional[str] = None,
-        **options,
-    ):
+    def deploy(self, streamlit: StreamlitEntityModel, replace: bool = False):
+        streamlit_id = streamlit.fqn.using_connection(self._conn)
+        if (
+            ObjectManager().object_exists(object_type="streamlit", fqn=streamlit_id)
+            and not replace
+        ):
+            raise ClickException(
+                f"Streamlit {streamlit.fqn} already exist. If you want to replace it use --replace flag."
+            )
+
         # 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
@@ -134,12 +146,9 @@ class StreamlitManager(SqlExecutionMixin):
             # 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:
@@ -169,10 +178,7 @@ class StreamlitManager(SqlExecutionMixin):
 
             self._put_streamlit_files(
                 root_location,
-                main_file,
-                environment_file,
-                pages_dir,
-                additional_source_files,
+                streamlit.artifacts,
             )
         else:
             """
@@ -182,31 +188,23 @@ class StreamlitManager(SqlExecutionMixin):
             """
             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(stage_name=stage_name)
+            cli_console.step(f"Creating {stage_name} stage")
+            stage_manager.create(fqn=stage_name)
 
             root_location = stage_manager.get_standard_stage_prefix(
                 f"{stage_name}/{streamlit_name_for_root_location}"
             )
 
-            self._put_streamlit_files(
-                root_location,
-                main_file,
-                environment_file,
-                pages_dir,
-                additional_source_files,
-            )
+            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)

snowflake/cli/_plugins/streamlit/plugin_spec.py

@@ -0,0 +1,30 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from snowflake.cli._plugins.streamlit import commands
+from snowflake.cli.api.plugins.command import (
+    SNOWCLI_ROOT_COMMAND_PATH,
+    CommandSpec,
+    CommandType,
+    plugin_hook_impl,
+)
+
+
+@plugin_hook_impl
+def command_spec():
+    return CommandSpec(
+        parent_command_path=SNOWCLI_ROOT_COMMAND_PATH,
+        command_type=CommandType.COMMAND_GROUP,
+        typer_instance=commands.app.create_instance(),
+    )

snowflake/cli/_plugins/streamlit/streamlit_entity.py

@@ -0,0 +1,12 @@
+from snowflake.cli._plugins.streamlit.streamlit_entity_model import (
+    StreamlitEntityModel,
+)
+from snowflake.cli.api.entities.common import EntityBase
+
+
+class StreamlitEntity(EntityBase[StreamlitEntityModel]):
+    """
+    A Streamlit app.
+    """
+
+    pass

snowflake/cli/_plugins/streamlit/streamlit_entity_model.py

@@ -0,0 +1,66 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from __future__ import annotations
+
+from pathlib import Path
+from typing import List, Literal, Optional
+
+from pydantic import Field, model_validator
+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, ImportsBaseModel):
+    type: Literal["streamlit"] = DiscriminatorField()  # noqa: A003
+    title: Optional[str] = Field(
+        title="Human-readable title for the Streamlit dashboard", default=None
+    )
+    comment: Optional[str] = Field(title="Comment for the Streamlit app", default=None)
+    query_warehouse: str = Field(
+        title="Snowflake warehouse to host the app", default=None
+    )
+    main_file: Optional[str] = Field(
+        title="Entrypoint file of the Streamlit app", default="streamlit_app.py"
+    )
+    pages_dir: Optional[str] = Field(title="Streamlit pages", default=None)
+    stage: Optional[str] = Field(
+        title="Stage in which the app’s artifacts will be stored", default="streamlit"
+    )
+    # Possibly can be PathMapping
+    artifacts: Optional[List[Path]] = Field(
+        title="List of files which should be deployed. Each file needs to exist locally. "
+        "Main file needs to be included in the artifacts.",
+        default=None,
+    )
+
+    @model_validator(mode="after")
+    def artifacts_must_exists(self):
+        if not self.artifacts:
+            return self
+
+        for artifact in self.artifacts:
+            if "*" in artifact.name:
+                continue
+            if not artifact.exists():
+                raise ValueError(
+                    f"Specified artifact {artifact} does not exist locally."
+                )
+
+        return self

snowflake/cli/_plugins/workspace/action_context.py

@@ -0,0 +1,18 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable, Optional
+
+from snowflake.cli.api.console.abc import AbstractConsole
+
+
+@dataclass
+class ActionContext:
+    """
+    An object that is passed to each action when called by WorkspaceManager
+    """
+
+    console: AbstractConsole
+    project_root: Path
+    default_role: str
+    default_warehouse: Optional[str]
+    get_entity: Callable

snowflake/cli/_plugins/workspace/commands.py

@@ -0,0 +1,306 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from textwrap import dedent
+from typing import List, Optional
+
+import typer
+from click import MissingParameter
+from snowflake.cli._plugins.nativeapp.artifacts import BundleMap
+from snowflake.cli._plugins.nativeapp.common_flags import (
+    ForceOption,
+    InteractiveOption,
+    ValidateOption,
+)
+from snowflake.cli._plugins.workspace.manager import WorkspaceManager
+from snowflake.cli.api.cli_global_context import get_cli_context
+from snowflake.cli.api.commands.decorators import with_project_definition
+from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
+from snowflake.cli.api.entities.common import EntityActions
+from snowflake.cli.api.exceptions import IncompatibleParametersError
+from snowflake.cli.api.output.types import MessageResult, QueryResult
+
+ws = SnowTyperFactory(
+    name="ws",
+    help="Deploy and interact with snowflake.yml-based entities.",
+    is_hidden=lambda: True,
+)
+log = logging.getLogger(__name__)
+
+
+@ws.command(requires_connection=True, hidden=True)
+@with_project_definition()
+def bundle(
+    entity_id: str = typer.Option(
+        help=f"""The ID of the entity you want to bundle.""",
+    ),
+    **options,
+):
+    """
+    Prepares a local folder with the configured artifacts of the specified entity.
+    """
+
+    cli_context = get_cli_context()
+    ws = WorkspaceManager(
+        project_definition=cli_context.project_definition,
+        project_root=cli_context.project_root,
+    )
+
+    bundle_map: BundleMap = ws.perform_action(entity_id, EntityActions.BUNDLE)
+    return MessageResult(f"Bundle generated at {bundle_map.deploy_root()}")
+
+
+@ws.command(requires_connection=True, hidden=True)
+@with_project_definition()
+def deploy(
+    entity_id: str = typer.Option(
+        help=f"""The ID of the entity you want to deploy.""",
+    ),
+    # TODO The following options should be generated automatically, depending on the specified entity type
+    prune: Optional[bool] = typer.Option(
+        default=None,
+        help=f"""Whether to delete specified files from the stage if they don't exist locally. If set, the command deletes files that exist in the stage, but not in the local filesystem. This option cannot be used when paths are specified.""",
+    ),
+    recursive: Optional[bool] = typer.Option(
+        None,
+        "--recursive/--no-recursive",
+        "-r",
+        help=f"""Whether to traverse and deploy files from subdirectories. If set, the command deploys all files and subdirectories; otherwise, only files in the current directory are deployed.""",
+    ),
+    paths: Optional[List[Path]] = typer.Argument(
+        default=None,
+        show_default=False,
+        help=dedent(
+            f"""
+            Paths, relative to the the project root, of files or directories you want to upload to a stage. If a file is
+            specified, it must match one of the artifacts src pattern entries in snowflake.yml. If a directory is
+            specified, it will be searched for subfolders or files to deploy based on artifacts src pattern entries. If
+            unspecified, the command syncs all local changes to the stage."""
+        ).strip(),
+    ),
+    from_release_directive: Optional[bool] = typer.Option(
+        False,
+        "--from-release-directive",
+        help=f"""Creates or upgrades an application object to the version and patch specified by the release directive applicable to your Snowflake account.
+        The command fails if no release directive exists for your Snowflake account for a given application package, which is determined from the project definition file. Default: unset.""",
+        is_flag=True,
+    ),
+    interactive: bool = InteractiveOption,
+    force: Optional[bool] = ForceOption,
+    validate: bool = ValidateOption,
+    **options,
+):
+    """
+    Deploys the specified entity.
+    """
+    if prune is None and recursive is None and not paths:
+        prune = True
+        recursive = True
+    else:
+        if prune is None:
+            prune = False
+        if recursive is None:
+            recursive = False
+
+    if paths and prune:
+        raise IncompatibleParametersError(["paths", "--prune"])
+
+    cli_context = get_cli_context()
+    ws = WorkspaceManager(
+        project_definition=cli_context.project_definition,
+        project_root=cli_context.project_root,
+    )
+
+    ws.perform_action(
+        entity_id,
+        EntityActions.DEPLOY,
+        prune=prune,
+        recursive=recursive,
+        paths=paths,
+        validate=validate,
+        from_release_directive=from_release_directive,
+        interactive=interactive,
+        force=force,
+    )
+    return MessageResult("Deployed successfully.")
+
+
+@ws.command(requires_connection=True, hidden=True)
+@with_project_definition()
+def drop(
+    entity_id: str = typer.Option(
+        help=f"""The ID of the entity you want to drop.""",
+    ),
+    # TODO The following options should be generated automatically, depending on the specified entity type
+    interactive: bool = InteractiveOption,
+    force: Optional[bool] = ForceOption,
+    cascade: Optional[bool] = typer.Option(
+        None,
+        help=f"""Whether to drop all application objects owned by the application within the account. Default: false.""",
+        show_default=False,
+    ),
+    **options,
+):
+    """
+    Drops the specified entity.
+    """
+    cli_context = get_cli_context()
+    ws = WorkspaceManager(
+        project_definition=cli_context.project_definition,
+        project_root=cli_context.project_root,
+    )
+
+    ws.perform_action(
+        entity_id,
+        EntityActions.DROP,
+        force_drop=force,
+        interactive=interactive,
+        cascade=cascade,
+    )
+
+
+@ws.command(requires_connection=True, hidden=True)
+@with_project_definition()
+def validate(
+    entity_id: str = typer.Option(
+        help=f"""The ID of the entity you want to validate.""",
+    ),
+    interactive: bool = InteractiveOption,
+    force: Optional[bool] = ForceOption,
+    **options,
+):
+    """Validates the specified entity."""
+    cli_context = get_cli_context()
+    ws = WorkspaceManager(
+        project_definition=cli_context.project_definition,
+        project_root=cli_context.project_root,
+    )
+
+    ws.perform_action(
+        entity_id,
+        EntityActions.VALIDATE,
+        interactive=interactive,
+        force=force,
+    )
+
+
+version = SnowTyperFactory(
+    name="version",
+    help="Manages versions for project entities.",
+)
+ws.add_typer(version)
+
+
+@version.command(name="list", requires_connection=True, hidden=True)
+@with_project_definition()
+def version_list(
+    entity_id: str = typer.Option(
+        help="The ID of the entity you want to list versions for.",
+    ),
+    **options,
+):
+    """Lists the versions of the specified entity."""
+    cli_context = get_cli_context()
+    ws = WorkspaceManager(
+        project_definition=cli_context.project_definition,
+        project_root=cli_context.project_root,
+    )
+    cursor = ws.perform_action(
+        entity_id,
+        EntityActions.VERSION_LIST,
+    )
+    return QueryResult(cursor)
+
+
+@version.command(name="create", requires_connection=True, hidden=True)
+@with_project_definition()
+def version_create(
+    entity_id: str = typer.Option(
+        help="The ID of the entity you want to create a version for.",
+    ),
+    version: Optional[str] = typer.Argument(
+        None,
+        help=f"""Version to define in your application package. If the version already exists, an auto-incremented patch is added to the version instead. Defaults to the version specified in the `manifest.yml` file.""",
+    ),
+    patch: Optional[int] = typer.Option(
+        None,
+        "--patch",
+        help=f"""The patch number you want to create for an existing version.
+        Defaults to undefined if it is not set, which means the Snowflake CLI either uses the patch specified in the `manifest.yml` file or automatically generates a new patch number.""",
+    ),
+    skip_git_check: Optional[bool] = typer.Option(
+        False,
+        "--skip-git-check",
+        help="When enabled, the Snowflake CLI skips checking if your project has any untracked or stages files in git. Default: unset.",
+        is_flag=True,
+    ),
+    interactive: bool = InteractiveOption,
+    force: Optional[bool] = ForceOption,
+    **options,
+):
+    """Creates a new version for the specified entity."""
+    if version is None and patch is not None:
+        raise MissingParameter("Cannot provide a patch without version!")
+
+    cli_context = get_cli_context()
+    ws = WorkspaceManager(
+        project_definition=cli_context.project_definition,
+        project_root=cli_context.project_root,
+    )
+    ws.perform_action(
+        entity_id,
+        EntityActions.VERSION_CREATE,
+        version=version,
+        patch=patch,
+        skip_git_check=skip_git_check,
+        interactive=interactive,
+        force=force,
+    )
+
+
+@version.command(name="drop", requires_connection=True, hidden=True)
+@with_project_definition()
+def version_drop(
+    entity_id: str = typer.Option(
+        help="The ID of the entity you want to create a version for.",
+    ),
+    version: Optional[str] = typer.Argument(
+        None,
+        help=f"""Version to define in your application package. If the version already exists, an auto-incremented patch is added to the version instead. Defaults to the version specified in the `manifest.yml` file.""",
+    ),
+    interactive: bool = InteractiveOption,
+    force: Optional[bool] = ForceOption,
+    **options,
+):
+    """
+    Drops a version defined for your entity. Versions can either be passed in as an argument to the command or read from the `manifest.yml` file.
+    Dropping patches is not allowed.
+    """
+
+    cli_context = get_cli_context()
+    ws = WorkspaceManager(
+        project_definition=cli_context.project_definition,
+        project_root=cli_context.project_root,
+    )
+    ws.perform_action(
+        entity_id,
+        EntityActions.VERSION_CREATE,
+        version=version,
+        interactive=interactive,
+        force=force,
+    )