Refactor containers commands

src/snowcli/cli/common/sql_execution.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from functools import cached_property
+from textwrap import dedent
 
 from snowcli.cli.common.snow_cli_global_context import snow_cli_global_context_manager
 
@@ -17,6 +18,6 @@ def _execute_template(self, template_name: str, payload: dict):
  return self._conn.run_sql(template_name, payload)
 
  def _execute_query(self, query: str):
- results = self._conn.ctx.execute_string(query)
+ results = self._conn.ctx.execute_string(dedent(query))
  *_, last_result = results
  return last_result

src/snowcli/cli/snowpark/__init__.py

@@ -4,9 +4,12 @@
 from snowcli.cli.snowpark.function import app as function_app
 from snowcli.cli.snowpark.package import app as package_app
 from snowcli.cli.snowpark.procedure import app as procedure_app
-from snowcli.cli.snowpark.cp import app as compute_pools_app, app_cp as cp_app
-from snowcli.cli.snowpark.services import app as services_app
-from snowcli.cli.snowpark.jobs import app as jobs_app
+from snowcli.cli.snowpark.compute_pool.commands import (
+ app as compute_pools_app,
+ app_cp as cp_app,
+)
+from snowcli.cli.snowpark.services.commands import app as services_app
+from snowcli.cli.snowpark.jobs.commands import app as jobs_app
 from snowcli.cli.snowpark.registry import app as registry_app
 
 app = typer.Typer(

src/snowcli/cli/snowpark/compute_pool/commands.py

@@ -0,0 +1,68 @@
+import typer
+
+from snowcli.cli.common.alias import build_alias
+from snowcli.cli.common.decorators import global_options
+from snowcli.cli.common.flags import DEFAULT_CONTEXT_SETTINGS
+from snowcli.cli.snowpark.compute_pool.manager import ComputePoolManager
+from snowcli.output.decorators import with_output
+
+
+app = typer.Typer(
+ context_settings=DEFAULT_CONTEXT_SETTINGS,
+ name="compute-pool",
+ help="Manage compute pools. You can also use cp as alias for this command",
+)
+
+
+@app.command()
+@with_output
+@global_options
+def create(
+ name: str = typer.Option(..., "--name", "-n", help="Compute pool name"),
+ num_instances: int = typer.Option(..., "--num", "-d", help="Number of instances"),
+ instance_family: str = typer.Option(..., "--family", "-f", help="Instance family"),
+ **options,
+):
+ """
+ Create compute pool
+ """
+ return ComputePoolManager().create(
+ pool_name=name, num_instances=num_instances, instance_family=instance_family
+ )
+
+
+@app.command()
+@with_output
+@global_options
+def list(**options):
+ """
+ List compute pools
+ """
+ return ComputePoolManager().show()
+
+
+@app.command()
+@with_output
+@global_options
+def drop(name: str = typer.Argument(..., help="Compute Pool Name"), **options):
+ """
+ Drop compute pool
+ """
+ return ComputePoolManager().drop(pool_name=name)
+
+
+@app.command()
+@with_output
+@global_options
+def stop(name: str = typer.Argument(..., help="Compute Pool Name"), **options):
+ """
+ Stop and delete all services running on Compute Pool
+ """
+ return ComputePoolManager().stop(pool_name=name)
+
+
+app_cp = build_alias(
+ app,
+ name="cp",
+ help_str="Manage compute pools. This is alias for compute-pool command",
+)

src/snowcli/cli/snowpark/compute_pool/manager.py

@@ -0,0 +1,25 @@
+from snowcli.cli.common.sql_execution import SqlExecutionMixin
+
+
+class ComputePoolManager(SqlExecutionMixin):
+ def create(self, pool_name: str, num_instances: int, instance_family: str):
+ return self._execute_query(
+ f"""\
+ CREATE COMPUTE POOL {pool_name}
+ MIN_NODES = {num_instances}
+ MAX_NODES = {num_instances}
+ INSTANCE_FAMILY = {instance_family};
+ """
+ )
+
+ def show(self):
+ return self._execute_query("show compute pools;")
+
+ def drop(
+ self,
+ pool_name: str,
+ ):
+ return self._execute_query(f"drop compute pool {pool_name};")
+
+ def stop(self, pool_name: str):
+ return self._execute_query(f"alter compute pool {pool_name} stop all services;")