SNOW-1011772: Suspend and resume commands for compute pools. (#740)

* Creating validation functions for object names.

* SNOW-1011772: Adding templates for suspend/resume

* SNOW-1011772: Adding 'resume' and 'suspend' commands for 'spcs compute-pool'.

* SNOW-1011772: Adding doc strings

* SNOW-1011772: Documentation fixes

* SNOW-1011772: Test fixes

* SNOW-1011772: Update release notes

* SNOW-1011772: Formatting

* SNOW-1011772: Updating compute pool callback to only allow a single valid identifier

RELEASE-NOTES.md

@@ -7,6 +7,7 @@
 * Introduced `snowflake.cli.api.console.cli_console` object with helper methods for intermediate output.
 * Added new convenience command `spcs image-registry url` to get the URL for your account image registry.
 * Added convenience function `spcs image-repository url <repo_name>`.
+* Added `suspend` and `resume` commands for `spcs compute-pool`.
 
 ## Fixes and improvements
 * Restricted permissions of automatically created files

src/snowflake/cli/api/project/util.py

@@ -12,8 +12,9 @@
 SINGLE_QUOTED_STRING_LITERAL_REGEX = r"'((?:\\.|''|[^'\n])+?)'"
 
 # See https://docs.snowflake.com/en/sql-reference/identifiers-syntax for identifier syntax
-UNQUOTED_IDENTIFIER_REGEX = r"(^[a-zA-Z_])([a-zA-Z0-9_$]{0,254})"
+UNQUOTED_IDENTIFIER_REGEX = r"([a-zA-Z_])([a-zA-Z0-9_$]{0,254})"
 QUOTED_IDENTIFIER_REGEX = r'"((""|[^"]){0,255})"'
+VALID_IDENTIFIER_REGEX = f"(?:{UNQUOTED_IDENTIFIER_REGEX}|{QUOTED_IDENTIFIER_REGEX})"
 
 
 def clean_identifier(input_: str):
@@ -47,6 +48,20 @@ def is_valid_identifier(identifier: str) -> bool:
     )
 
 
+def is_valid_object_name(name: str, max_depth=2) -> bool:
+    """
+    Determines whether the given identifier is a valid object name in the form <name>, <schema>.<name>, or <database>.<schema>.<name>.
+    Max_depth determines how many valid identifiers are allowed. For example, account level objects would have a max depth of 0
+    because they cannot be qualified by a database or schema, just the single identifier.
+    """
+    if max_depth < 0:
+        raise ValueError("max_depth must be non-negative")
+    pattern = (
+        rf"{VALID_IDENTIFIER_REGEX}(?:\.{VALID_IDENTIFIER_REGEX}){{0,{max_depth}}}"
+    )
+    return re.fullmatch(pattern, name) is not None
+
+
 def to_identifier(name: str) -> str:
     """
     Converts a name to a valid Snowflake identifier. If the name is already a valid

src/snowflake/cli/plugins/spcs/compute_pool/commands.py

@@ -1,12 +1,14 @@
 from typing import Optional
 
 import typer
+from click import ClickException
 from snowflake.cli.api.commands.decorators import (
     global_options_with_connection,
     with_output,
 )
 from snowflake.cli.api.commands.flags import DEFAULT_CONTEXT_SETTINGS
 from snowflake.cli.api.output.types import CommandResult, SingleQueryResult
+from snowflake.cli.api.project.util import is_valid_object_name
 from snowflake.cli.plugins.object.common import comment_option
 from snowflake.cli.plugins.spcs.common import validate_and_set_instances
 from snowflake.cli.plugins.spcs.compute_pool.manager import ComputePoolManager
@@ -18,11 +20,25 @@
 )
 
 
+def _compute_pool_name_callback(name: str) -> str:
+    """
+    Verifies that compute pool name is a single valid identifier.
+    """
+    if not is_valid_object_name(name, 0):
+        raise ClickException(f"{name} is not a valid compute pool name.")
+    return name
+
+
+ComputePoolNameArgument = typer.Argument(
+    ..., help="Name of the compute pool.", callback=_compute_pool_name_callback
+)
+
+
 @app.command()
 @with_output
 @global_options_with_connection
 def create(
-    name: str = typer.Argument(..., help="Name of the compute pool."),
+    name: str = ComputePoolNameArgument,
     min_nodes: int = typer.Option(
         1, "--min-nodes", help="Minimum number of nodes for the compute pool."
     ),
@@ -72,11 +88,29 @@ def create(
 @app.command("stop-all")
 @with_output
 @global_options_with_connection
-def stop_all(
-    name: str = typer.Argument(..., help="Name of the compute pool."), **options
-) -> CommandResult:
+def stop_all(name: str = ComputePoolNameArgument, **options) -> CommandResult:
     """
-    Stops a compute pool and deletes all services running on the pool.
+    Deletes all services running on the compute pool.
     """
     cursor = ComputePoolManager().stop(pool_name=name)
     return SingleQueryResult(cursor)
+
+
+@app.command()
+@with_output
+@global_options_with_connection
+def suspend(name: str = ComputePoolNameArgument, **options) -> CommandResult:
+    """
+    Suspends the compute pool by suspending all currently running services and then releasing compute pool nodes.
+    """
+    return SingleQueryResult(ComputePoolManager().suspend(name))
+
+
+@app.command()
+@with_output
+@global_options_with_connection
+def resume(name: str = ComputePoolNameArgument, **options) -> CommandResult:
+    """
+    Resumes the compute pool from SUSPENDED state.
+    """
+    return SingleQueryResult(ComputePoolManager().resume(name))

src/snowflake/cli/plugins/spcs/compute_pool/manager.py

@@ -31,4 +31,10 @@ def create(
         return self._execute_query(strip_empty_lines(query))
 
     def stop(self, pool_name: str) -> SnowflakeCursor:
-        return self._execute_query(f"alter compute pool {pool_name} stop all;")
+        return self._execute_query(f"alter compute pool {pool_name} stop all")
+
+    def suspend(self, pool_name: str) -> SnowflakeCursor:
+        return self._execute_query(f"alter compute pool {pool_name} suspend")
+
+    def resume(self, pool_name: str) -> SnowflakeCursor:
+        return self._execute_query(f"alter compute pool {pool_name} resume")