Output types (snowflakedb#350)

src/snowcli/app/cli_app.py

@@ -12,7 +12,8 @@
 from snowcli.config import config_init, cli_config
 from snowcli.app.dev.docs.generator import generate_docs
 from snowcli.output.formats import OutputFormat
-from snowcli.output.printing import OutputData
+from snowcli.output.printing import print_result
+from snowcli.output.types import CollectionResult
 from snowcli.app.dev.pycharm_remote_debug import (
  setup_pycharm_remote_debugger_if_provided,
 )
@@ -36,13 +37,13 @@ def _version_callback(value: bool):
 
 def _info_callback(value: bool):
  if value:
- OutputData.from_list(
+ result = CollectionResult(
  [
  {"key": "version", "value": __about__.VERSION},
  {"key": "default_config_file_path", "value": cli_config.file_path},
  ],
-  format_=OutputFormat.JSON,
- ).print()
+ )
+ print_result(result, output_format=OutputFormat.JSON)
  raise typer.Exit()
 
 

src/snowcli/cli/connection/commands.py

@@ -11,7 +11,7 @@
 from snowcli.cli.common.flags import DEFAULT_CONTEXT_SETTINGS, ConnectionOption
 from snowcli.config import cli_config
 from snowcli.output.decorators import with_output
-from snowcli.output.printing import OutputData
+from snowcli.output.types import CollectionResult, CommandResult, MessageResult
 from snowcli.snow_connector import connect_to_snowflake
 
 app = typer.Typer(
@@ -41,7 +41,7 @@ def _mask_password(connection_params: dict):
 @app.command(name="list")
 @with_output
 @global_options
-def list_connections(**options) -> OutputData:
+def list_connections(**options) -> CommandResult:
  """
  List configured connections.
  """
@@ -50,7 +50,7 @@ def list_connections(**options) -> OutputData:
  {"connection_name": k, "parameters": _mask_password(v)}
  for k, v in connections.items()
  )
- return OutputData(stream=result)
+ return CollectionResult(result)
 
 
 def require_integer(field_name: str):
@@ -159,7 +159,7 @@ def add(
  prompt="Snowflake region",
  help="Region name if not the default Snowflake deployment.",
  ),
-) -> OutputData:
+) -> CommandResult:
  """Add connection to configuration file."""
  connection_entry = {
  "account": account,
@@ -180,16 +180,16 @@ def add(
  except KeyAlreadyPresent:
  raise ClickException(f"Connection {connection_name} already exists")
 
- return OutputData.from_string(
+ return MessageResult(
  f"Wrote new connection {connection_name} to {cli_config.file_path}"
  )
 
 
 @app.command()
 @with_output
-def test(connection: str = ConnectionOption) -> OutputData:
+def test(connection: str = ConnectionOption) -> CommandResult:
  """
  Tests connection to Snowflake.
  """
  connect_to_snowflake(connection_name=connection)
- return OutputData.from_string("OK")
+ return MessageResult("OK")

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

@@ -5,7 +5,7 @@
 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
-from snowcli.output.printing import OutputData
+from snowcli.output.types import SingleQueryResult, QueryResult, CommandResult
 
 app = typer.Typer(
  context_settings=DEFAULT_CONTEXT_SETTINGS,
@@ -22,51 +22,51 @@ def create(
  num_instances: int = typer.Option(..., "--num", "-d", help="Number of instances"),
  instance_family: str = typer.Option(..., "--family", "-f", help="Instance family"),
  **options,
-) -> OutputData:
+) -> CommandResult:
  """
  Create compute pool
  """
  cursor = ComputePoolManager().create(
  pool_name=name, num_instances=num_instances, instance_family=instance_family
  )
- return OutputData.from_cursor(cursor)
+ return SingleQueryResult(cursor)
 
 
 @app.command()
 @with_output
 @global_options_with_connection
-def list(**options) -> OutputData:
+def list(**options) -> CommandResult:
  """
  List compute pools
  """
  cursor = ComputePoolManager().show()
- return OutputData.from_cursor(cursor)
+ return QueryResult(cursor)
 
 
 @app.command()
 @with_output
 @global_options_with_connection
 def drop(
  name: str = typer.Argument(..., help="Compute Pool Name"), **options
-) -> OutputData:
+) -> CommandResult:
  """
  Drop compute pool
  """
  cursor = ComputePoolManager().drop(pool_name=name)
- return OutputData.from_cursor(cursor)
+ return SingleQueryResult(cursor)
 
 
 @app.command()
 @with_output
 @global_options_with_connection
 def stop(
  name: str = typer.Argument(..., help="Compute Pool Name"), **options
-) -> OutputData:
+) -> CommandResult:
  """
  Stop and delete all services running on Compute Pool
  """
  cursor = ComputePoolManager().stop(pool_name=name)
- return OutputData.from_cursor(cursor)
+ return SingleQueryResult(cursor)
 
 
 app_cp = build_alias(

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

@@ -18,7 +18,12 @@
 )
 from snowcli.cli.stage.manager import StageManager
 from snowcli.output.decorators import with_output
-from snowcli.output.printing import OutputData
+from snowcli.output.types import (
+ MessageResult,
+ SingleQueryResult,
+ QueryResult,
+ CommandResult,
+)
 from snowcli.utils import (
  prepare_app_zip,
  get_snowflake_packages,
@@ -71,7 +76,7 @@ def function_init():
  Initialize this directory with a sample set of files to create a function.
  """
  create_project_template("default_function")
- return OutputData.from_string("Done")
+ return MessageResult("Done")
 
 
 @app.command("create")
@@ -104,7 +109,7 @@ def function_create(
  help="Replace if existing function",
  ),
  **options,
-) -> OutputData:
+) -> CommandResult:
  """Creates a python UDF/UDTF using local artifact."""
  snowpark_package(
  pypi_download, # type: ignore[arg-type]
@@ -133,7 +138,7 @@ def function_create(
  packages=packages,
  overwrite=overwrite,
  )
- return OutputData.from_cursor(cursor)
+ return SingleQueryResult(cursor)
 
 
 def upload_snowpark_artifact(
@@ -182,7 +187,7 @@ def function_update(
  help="Replace function, even if no detected changes to metadata",
  ),
  **options,
-) -> OutputData:
+) -> CommandResult:
  """Updates an existing python UDF/UDTF using local artifact."""
  snowpark_package(
  pypi_download, # type: ignore[arg-type]
@@ -241,9 +246,9 @@ def function_update(
  packages=packages,
  overwrite=True,
  )
- return OutputData.from_cursor(cursor)
+ return SingleQueryResult(cursor)
 
- return OutputData.from_string("No packages to update. Deployment complete!")
+ return MessageResult("No packages to update. Deployment complete!")
 
 
 @app.command("package")
@@ -252,14 +257,14 @@ def function_package(
  pypi_download: str = PyPiDownloadOption,
  check_anaconda_for_pypi_deps: bool = CheckAnacondaForPyPiDependancies,
  package_native_libraries: str = PackageNativeLibrariesOption,
-) -> OutputData:
+) -> CommandResult:
  """Packages function code into zip file."""
  snowpark_package(
  pypi_download, # type: ignore[arg-type]
  check_anaconda_for_pypi_deps,
  package_native_libraries, # type: ignore[arg-type]
  )
- return OutputData.from_string("Done")
+ return MessageResult("Done")
 
 
 @app.command("execute")
@@ -273,10 +278,10 @@ def function_execute(
  help="Function with inputs. E.g. 'hello(int, string)'",
  ),
  **options,
-) -> OutputData:
+) -> CommandResult:
  """Executes a Snowflake function."""
  cursor = FunctionManager().execute(expression=function)
- return OutputData.from_cursor(cursor)
+ return SingleQueryResult(cursor)
 
 
 @app.command("describe")
@@ -292,14 +297,14 @@ def function_describe(
  help="Function signature with inputs. E.g. 'hello(int, string)'",
  ),
  **options,
-) -> OutputData:
+) -> CommandResult:
  """Describes a Snowflake function."""
  cursor = FunctionManager().describe(
  identifier=FunctionManager.identifier(
  name=name, signature=input_parameters, name_and_signature=function
  )
  )
- return OutputData.from_cursor(cursor)
+ return QueryResult(cursor)
 
 
 @app.command("list")
@@ -313,10 +318,10 @@ def function_list(
  help='Filter functions by name - e.g. "hello%"',
  ),
  **options,
-) -> OutputData:
+) -> CommandResult:
  """Lists Snowflake functions."""
  cursor = FunctionManager().show(like=like)
- return OutputData.from_cursor(cursor)
+ return QueryResult(cursor)
 
 
 @app.command("drop")
@@ -332,11 +337,11 @@ def function_drop(
  help="Function signature with inputs. E.g. 'hello(int, string)'",
  ),
  **options,
-) -> OutputData:
+) -> CommandResult:
  """Drops a Snowflake function."""
  cursor = FunctionManager().drop(
  identifier=FunctionManager.identifier(
  name=name, signature=input_parameters, name_and_signature=signature
  )
  )
- return OutputData.from_cursor(cursor)
+ return SingleQueryResult(cursor)

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

@@ -9,7 +9,8 @@
 from snowcli.cli.snowpark.jobs.manager import JobManager
 from snowcli.cli.stage.manager import StageManager
 from snowcli.output.decorators import with_output
-from snowcli.output.printing import OutputData
+
+from snowcli.output.types import SingleQueryResult, CommandResult
 
 app = typer.Typer(
  context_settings=DEFAULT_CONTEXT_SETTINGS, name="jobs", help="Manage jobs"
@@ -32,7 +33,7 @@ def create(
  ),
  stage: str = typer.Option("SOURCE_STAGE", "--stage", "-l", help="Stage name"),
  **options,
-) -> OutputData:
+) -> CommandResult:
  """
  Create Job
  """
@@ -43,18 +44,18 @@ def create(
  cursor = JobManager().create(
  compute_pool=compute_pool, spec_path=spec_path, stage=stage
  )
- return OutputData.from_cursor(cursor)
+ return SingleQueryResult(cursor)
 
 
 @app.command()
 @with_output
 @global_options_with_connection
-def desc(id: str = typer.Argument(..., help="Job id"), **options) -> OutputData:
+def desc(id: str = typer.Argument(..., help="Job id"), **options) -> CommandResult:
  """
  Desc Service
  """
  cursor = JobManager().desc(job_name=id)
- return OutputData.from_cursor(cursor)
+ return SingleQueryResult(cursor)
 
 
 @app.command()
@@ -78,20 +79,20 @@ def logs(
 @app.command()
 @with_output
 @global_options_with_connection
-def status(id: str = typer.Argument(..., help="Job id"), **options) -> OutputData:
+def status(id: str = typer.Argument(..., help="Job id"), **options) -> CommandResult:
  """
  Returns status of a job.
  """
  cursor = JobManager().status(job_name=id)
- return OutputData.from_cursor(cursor)
+ return SingleQueryResult(cursor)
 
 
 @app.command()
 @with_output
 @global_options_with_connection
-def drop(id: str = typer.Argument(..., help="Job id"), **options) -> OutputData:
+def drop(id: str = typer.Argument(..., help="Job id"), **options) -> CommandResult:
  """
  Drop Service
  """
  cursor = JobManager().drop(job_name=id)
- return OutputData.from_cursor(cursor)
+ return SingleQueryResult(cursor)