Skip to content

Commit

Permalink
Added snow spcs compute-pool deploy command
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfc-gh-astus
sfc-gh-astus committed Jan 17, 2025
1 parent b39d50b commit f689459
Show file tree
Hide file tree
Showing 13 changed files with 555 additions and 32 deletions.
2 changes: 2 additions & 0 deletions RELEASE-NOTES.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,8 @@
* Use `--all` to fetch all columns.
* Added support for glob pattern (except `**`) in artifact paths in snowflake.yml for Streamlit.
* Added support for glob pattern (except `**`) in artifact paths in snowflake.yml for Snowpark, requires ENABLE_SNOWPARK_GLOB_SUPPORT feature flag.
* Added `--replace` flag to `snow spcs compute-pool create` command.
* Added command `snow spcs compute-pool deploy`.

## Fixes and improvements

Expand Down
59 changes: 56 additions & 3 deletions src/snowflake/cli/_plugins/spcs/compute_pool/commands.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,28 +14,40 @@

from __future__ import annotations

from typing import Optional
from typing import Dict, Optional

import typer
from click import ClickException
from click import ClickException, UsageError
from snowflake.cli._plugins.object.command_aliases import (
add_object_command_aliases,
)
from snowflake.cli._plugins.object.common import CommentOption
from snowflake.cli._plugins.spcs.common import (
validate_and_set_instances,
)
from snowflake.cli._plugins.spcs.compute_pool.compute_pool_entity_model import (
ComputePoolEntityModel,
)
from snowflake.cli._plugins.spcs.compute_pool.manager import ComputePoolManager
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.flags import (
IfNotExistsOption,
OverrideableOption,
ReplaceOption,
entity_argument,
identifier_argument,
like_option,
)
from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
from snowflake.cli.api.constants import ObjectType
from snowflake.cli.api.exceptions import NoProjectDefinitionError
from snowflake.cli.api.identifiers import FQN
from snowflake.cli.api.output.types import CommandResult, SingleQueryResult
from snowflake.cli.api.output.types import (
CommandResult,
MessageResult,
SingleQueryResult,
)
from snowflake.cli.api.project.util import is_valid_object_name

app = SnowTyperFactory(
Expand Down Expand Up @@ -124,6 +136,7 @@ def create(
),
auto_suspend_secs: int = AutoSuspendSecsOption(),
comment: Optional[str] = CommentOption(help=_COMMENT_HELP),
replace: bool = ReplaceOption(),
if_not_exists: bool = IfNotExistsOption(),
**options,
) -> CommandResult:
Expand All @@ -141,10 +154,50 @@ def create(
auto_suspend_secs=auto_suspend_secs,
comment=comment,
if_not_exists=if_not_exists,
replace=replace,
)
return SingleQueryResult(cursor)


@app.command("deploy", requires_connection=True)
@with_project_definition()
def deploy(
replace: bool = ReplaceOption(
help="Replace the compute-pool if it already exists."
),
entity_id: str = entity_argument("compute-pool"),
**options,
):
"""
Deploys a compute pool from the project definition file.
"""
cli_context = get_cli_context()
pd = cli_context.project_definition
compute_pools: Dict[str, ComputePoolEntityModel] = pd.get_entities_by_type(
entity_type="compute-pool"
)

if not compute_pools:
raise NoProjectDefinitionError(
project_type="compute-pool", project_root=cli_context.project_root
)

if entity_id and entity_id not in compute_pools:
raise UsageError(f"No '{entity_id}' entity in project definition file.")

if len(compute_pools.keys()) == 1:
entity_id = list(compute_pools.keys())[0]

if entity_id is None:
raise UsageError(
"Multiple compute-pools found. Please provide entity id for the operation."
)

ComputePoolManager().deploy(compute_pool=compute_pools[entity_id], replace=replace)

return MessageResult(f"Compute pool '{entity_id}' successfully deployed.")


@app.command("stop-all", requires_connection=True)
def stop_all(name: FQN = ComputePoolNameArgument, **options) -> CommandResult:
"""
Expand Down
8 changes: 8 additions & 0 deletions src/snowflake/cli/_plugins/spcs/compute_pool/compute_pool_entity.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
from snowflake.cli._plugins.spcs.compute_pool.compute_pool_entity_model import (
ComputePoolEntityModel,
)
from snowflake.cli.api.entities.common import EntityBase


class ComputePoolEntity(EntityBase[ComputePoolEntityModel]):
pass
23 changes: 23 additions & 0 deletions src/snowflake/cli/_plugins/spcs/compute_pool/compute_pool_entity_model.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,23 @@
from typing import Literal, Optional

from pydantic import Field
from snowflake.cli.api.project.schemas.entities.common import EntityModelBase
from snowflake.cli.api.project.schemas.updatable_model import DiscriminatorField


class ComputePoolEntityModel(EntityModelBase):
type: Literal["compute-pool"] = DiscriminatorField() # noqa: A003
min_nodes: Optional[int] = Field(title="Minimum number of nodes", default=None)
max_nodes: Optional[int] = Field(title="Maximum number of nodes", default=None)
instance_family: str = Field(title="Name of the instance family", default=None)
auto_resume: Optional[bool] = Field(
title="The compute pool will automatically resume when a service or job is submitted to it",
default=True,
)
initially_suspended: Optional[bool] = Field(
title="Starts the compute pool in a suspended state", default=False
)
auto_suspend_seconds: Optional[int] = Field(
title="Number of seconds of inactivity after which you want Snowflake to automatically suspend the compute pool",
default=3600,
)
35 changes: 34 additions & 1 deletion src/snowflake/cli/_plugins/spcs/compute_pool/manager.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,12 +16,17 @@

from typing import List, Optional

from snowflake.cli._plugins.object.manager import ObjectManager
from snowflake.cli._plugins.spcs.common import (
NoPropertiesProvidedError,
handle_object_already_exists,
strip_empty_lines,
)
from snowflake.cli._plugins.spcs.compute_pool.compute_pool_entity_model import (
ComputePoolEntityModel,
)
from snowflake.cli.api.constants import ObjectType
from snowflake.cli.api.identifiers import FQN
from snowflake.cli.api.sql_execution import SqlExecutionMixin
from snowflake.connector.cursor import SnowflakeCursor
from snowflake.connector.errors import ProgrammingError
Expand All @@ -39,7 +44,17 @@ def create(
auto_suspend_secs: int,
comment: Optional[str],
if_not_exists: bool,
replace: bool,
) -> SnowflakeCursor:

if replace:
object_manager = ObjectManager()
object_type = str(ObjectType.COMPUTE_POOL)
entity_id_fqn = FQN.from_string(pool_name)
if object_manager.object_exists(object_type=object_type, fqn=entity_id_fqn):
self.stop(pool_name)
object_manager.drop(object_type=object_type, fqn=entity_id_fqn)

create_statement = "CREATE COMPUTE POOL"
if if_not_exists:
create_statement = f"{create_statement} IF NOT EXISTS"
Expand All @@ -58,7 +73,25 @@ def create(
try:
return self.execute_query(strip_empty_lines(query))
except ProgrammingError as e:
handle_object_already_exists(e, ObjectType.COMPUTE_POOL, pool_name)
handle_object_already_exists(
e, ObjectType.COMPUTE_POOL, pool_name, replace_available=True
)

def deploy(
self, compute_pool: ComputePoolEntityModel, replace: bool
) -> SnowflakeCursor:
return self.create(
pool_name=compute_pool.entity_id,
min_nodes=compute_pool.min_nodes,
max_nodes=compute_pool.max_nodes,
instance_family=compute_pool.instance_family,
auto_resume=compute_pool.auto_resume,
initially_suspended=compute_pool.initially_suspended,
auto_suspend_secs=compute_pool.auto_suspend_seconds,
comment=None,
if_not_exists=False,
replace=replace,
)

def stop(self, pool_name: str) -> SnowflakeCursor:
return self.execute_query(f"alter compute pool {pool_name} stop all")
Expand Down
8 changes: 8 additions & 0 deletions src/snowflake/cli/api/project/schemas/entities/entities.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,12 @@
FunctionEntityModel,
ProcedureEntityModel,
)
from snowflake.cli._plugins.spcs.compute_pool.compute_pool_entity import (
ComputePoolEntity,
)
from snowflake.cli._plugins.spcs.compute_pool.compute_pool_entity_model import (
ComputePoolEntityModel,
)
from snowflake.cli._plugins.streamlit.streamlit_entity import StreamlitEntity
from snowflake.cli._plugins.streamlit.streamlit_entity_model import (
StreamlitEntityModel,
Expand All @@ -43,13 +49,15 @@
StreamlitEntity,
ProcedureEntity,
FunctionEntity,
ComputePoolEntity,
]
EntityModel = Union[
ApplicationEntityModel,
ApplicationPackageEntityModel,
StreamlitEntityModel,
FunctionEntityModel,
ProcedureEntityModel,
ComputePoolEntityModel,
]

ALL_ENTITIES: List[Entity] = [*get_args(Entity)]
Expand Down
99 changes: 99 additions & 0 deletions tests/__snapshots__/test_help_messages.ambr
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5662,6 +5662,9 @@
| --comment TEXT Comment for the |
| compute pool. |
| [default: None] |
| --replace Replace this |
| object if it |
| already exists. |
| --if-not-exists Only apply this |
| operation if the |
| specified object |
Expand Down Expand Up @@ -5744,6 +5747,100 @@
+------------------------------------------------------------------------------+


'''
# ---
# name: test_help_messages[spcs.compute-pool.deploy]
'''

Usage: default spcs compute-pool deploy [OPTIONS] [ENTITY_ID]

Deploys a compute pool from the project definition file.

+- Arguments ------------------------------------------------------------------+
| entity_id [ENTITY_ID] ID of compute-pool entity. |
| [default: None] |
+------------------------------------------------------------------------------+
+- Options --------------------------------------------------------------------+
| --replace Replace the compute-pool if it already exists. |
| --project -p TEXT Path where Snowflake project resides. Defaults to |
| current working directory. |
| --env TEXT String in format of key=value. Overrides variables |
| from env section used for templates. |
| --help -h Show this message and exit. |
+------------------------------------------------------------------------------+
+- Connection configuration ---------------------------------------------------+
| --connection,--environment -c TEXT Name of the connection, as |
| defined in your config.toml |
| file. Default: default. |
| --host TEXT Host address for the |
| connection. Overrides the |
| value specified for the |
| connection. |
| --port INTEGER Port for the connection. |
| Overrides the value |
| specified for the |
| connection. |
| --account,--accountname TEXT Name assigned to your |
| Snowflake account. Overrides |
| the value specified for the |
| connection. |
| --user,--username TEXT Username to connect to |
| Snowflake. Overrides the |
| value specified for the |
| connection. |
| --password TEXT Snowflake password. |
| Overrides the value |
| specified for the |
| connection. |
| --authenticator TEXT Snowflake authenticator. |
| Overrides the value |
| specified for the |
| connection. |
| --private-key-file,--private… TEXT Snowflake private key file |
| path. Overrides the value |
| specified for the |
| connection. |
| --token-file-path TEXT Path to file with an OAuth |
| token that should be used |
| when connecting to Snowflake |
| --database,--dbname TEXT Database to use. Overrides |
| the value specified for the |
| connection. |
| --schema,--schemaname TEXT Database schema to use. |
| Overrides the value |
| specified for the |
| connection. |
| --role,--rolename TEXT Role to use. Overrides the |
| value specified for the |
| connection. |
| --warehouse TEXT Warehouse to use. Overrides |
| the value specified for the |
| connection. |
| --temporary-connection -x Uses connection defined with |
| command line parameters, |
| instead of one defined in |
| config |
| --mfa-passcode TEXT Token to use for |
| multi-factor authentication |
| (MFA) |
| --enable-diag Run Python connector |
| diagnostic test |
| --diag-log-path TEXT Diagnostic report path |
| --diag-allowlist-path TEXT Diagnostic report path to |
| optional allowlist |
+------------------------------------------------------------------------------+
+- Global configuration -------------------------------------------------------+
| --format [TABLE|JSON] Specifies the output format. |
| [default: TABLE] |
| --verbose -v Displays log entries for log levels info |
| and higher. |
| --debug Displays log entries for log levels debug |
| and higher; debug logs contain additional |
| information. |
| --silent Turns off intermediate output to console. |
+------------------------------------------------------------------------------+


'''
# ---
# name: test_help_messages[spcs.compute-pool.describe]
Expand Down Expand Up @@ -6603,6 +6700,7 @@
+------------------------------------------------------------------------------+
+- Commands -------------------------------------------------------------------+
| create Creates a new compute pool. |
| deploy Deploys a compute pool from the project definition file. |
| describe Provides description of compute pool. |
| drop Drops compute pool with given name. |
| list Lists all available compute pools. |
Expand Down Expand Up @@ -11113,6 +11211,7 @@
+------------------------------------------------------------------------------+
+- Commands -------------------------------------------------------------------+
| create Creates a new compute pool. |
| deploy Deploys a compute pool from the project definition file. |
| describe Provides description of compute pool. |
| drop Drops compute pool with given name. |
| list Lists all available compute pools. |
Expand Down
Loading

0 comments on commit f689459

Please sign in to comment.