snowflake-cli-labs 2.8.0rc0__py3-none-any.whl → 3.0.0rc0__py3-none-any.whl

snowflake/cli/{plugins → _plugins}/snowpark/zipper.py

@@ -17,7 +17,7 @@ from __future__ import annotations
 import fnmatch
 import logging
 from pathlib import Path
-from typing import Iterator, Literal
+from typing import Dict, List, Literal
 from zipfile import ZIP_DEFLATED, ZipFile
 
 log = logging.getLogger(__name__)
@@ -59,16 +59,24 @@ def add_file_to_existing_zip(zip_file: str, file: str):
 
 
 def zip_dir(
-    source: Path, dest_zip: Path, mode: Literal["r", "w", "x", "a"] = "w"
+    source: Path | List[Path],
+    dest_zip: Path,
+    mode: Literal["r", "w", "x", "a"] = "w",
 ) -> None:
-    files_to_pack: Iterator[Path] = filter(
-        _to_be_zipped, map(lambda f: f.absolute(), source.glob("**/*"))
-    )
+
+    if isinstance(source, Path):
+        source = [source]
+
+    files_to_pack: Dict[Path, List[Path]] = {
+        src: list(filter(_to_be_zipped, (f.absolute() for f in src.glob("**/*"))))
+        for src in source
+    }
 
     with ZipFile(dest_zip, mode, ZIP_DEFLATED, allowZip64=True) as package_zip:
-        for file in files_to_pack:
-            log.debug("Adding %s to %s", file, dest_zip)
-            package_zip.write(file, arcname=file.relative_to(source.absolute()))
+        for src, files in files_to_pack.items():
+            for file in files:
+                log.debug("Adding %s to %s", file, dest_zip)
+                package_zip.write(file, arcname=file.relative_to(src))
 
 
 def _to_be_zipped(file: Path) -> bool:

snowflake/cli/{plugins → _plugins}/spcs/__init__.py

@@ -12,16 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
-from snowflake.cli.plugins.spcs.compute_pool.commands import (
+from snowflake.cli._plugins.spcs.compute_pool.commands import (
     app as compute_pools_app,
 )
-from snowflake.cli.plugins.spcs.image_registry.commands import app as registry_app
-from snowflake.cli.plugins.spcs.image_repository.commands import (
+from snowflake.cli._plugins.spcs.image_registry.commands import app as registry_app
+from snowflake.cli._plugins.spcs.image_repository.commands import (
     app as image_repository_app,
 )
-from snowflake.cli.plugins.spcs.jobs.commands import app as jobs_app
-from snowflake.cli.plugins.spcs.services.commands import app as services_app
+from snowflake.cli._plugins.spcs.services.commands import app as services_app
+from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
 
 app = SnowTyperFactory(
     name="spcs",
@@ -30,6 +29,5 @@ app = SnowTyperFactory(
 
 app.add_typer(compute_pools_app)
 app.add_typer(services_app)
-app.add_typer(jobs_app)
 app.add_typer(registry_app)
 app.add_typer(image_repository_app)

snowflake/cli/{plugins → _plugins}/spcs/compute_pool/commands.py

@@ -18,23 +18,25 @@ from typing import Optional
 
 import typer
 from click import ClickException
+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.manager import ComputePoolManager
 from snowflake.cli.api.commands.flags import (
     IfNotExistsOption,
     OverrideableOption,
+    identifier_argument,
     like_option,
 )
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
 from snowflake.cli.api.constants import ObjectType
+from snowflake.cli.api.identifiers import FQN
 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.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.manager import ComputePoolManager
 
 app = SnowTyperFactory(
     name="compute-pool",
@@ -43,22 +45,21 @@ app = SnowTyperFactory(
 )
 
 
-def _compute_pool_name_callback(name: str) -> str:
+def _compute_pool_name_callback(name: FQN) -> FQN:
     """
     Verifies that compute pool name is a single valid identifier.
     """
-    if not is_valid_object_name(name, max_depth=0, allow_quoted=False):
+    if not is_valid_object_name(name.identifier, max_depth=0, allow_quoted=False):
         raise ClickException(
             f"'{name}' is not a valid compute pool name. Note that compute pool names must be unquoted identifiers."
         )
     return name
 
 
-ComputePoolNameArgument = typer.Argument(
-    ...,
-    help="Name of the compute pool.",
+ComputePoolNameArgument = identifier_argument(
+    sf_object="compute pool",
+    example="my_compute_pool",
     callback=_compute_pool_name_callback,
-    show_default=False,
 )
 
 
@@ -106,7 +107,7 @@ add_object_command_aliases(
 
 @app.command(requires_connection=True)
 def create(
-    name: str = ComputePoolNameArgument,
+    name: FQN = ComputePoolNameArgument,
     instance_family: str = typer.Option(
         ...,
         "--family",
@@ -131,7 +132,7 @@ def create(
     """
     max_nodes = validate_and_set_instances(min_nodes, max_nodes, "nodes")
     cursor = ComputePoolManager().create(
-        pool_name=name,
+        pool_name=name.identifier,
         min_nodes=min_nodes,
         max_nodes=max_nodes,
         instance_family=instance_family,
@@ -145,33 +146,33 @@ def create(
 
 
 @app.command("stop-all", requires_connection=True)
-def stop_all(name: str = ComputePoolNameArgument, **options) -> CommandResult:
+def stop_all(name: FQN = ComputePoolNameArgument, **options) -> CommandResult:
     """
     Deletes all services running on the compute pool.
     """
-    cursor = ComputePoolManager().stop(pool_name=name)
+    cursor = ComputePoolManager().stop(pool_name=name.identifier)
     return SingleQueryResult(cursor)
 
 
 @app.command(requires_connection=True)
-def suspend(name: str = ComputePoolNameArgument, **options) -> CommandResult:
+def suspend(name: FQN = ComputePoolNameArgument, **options) -> CommandResult:
     """
     Suspends the compute pool by suspending all currently running services and then releasing compute pool nodes.
     """
-    return SingleQueryResult(ComputePoolManager().suspend(name))
+    return SingleQueryResult(ComputePoolManager().suspend(name.identifier))
 
 
 @app.command(requires_connection=True)
-def resume(name: str = ComputePoolNameArgument, **options) -> CommandResult:
+def resume(name: FQN = ComputePoolNameArgument, **options) -> CommandResult:
     """
     Resumes the compute pool from a SUSPENDED state.
     """
-    return SingleQueryResult(ComputePoolManager().resume(name))
+    return SingleQueryResult(ComputePoolManager().resume(name.identifier))
 
 
 @app.command("set", requires_connection=True)
 def set_property(
-    name: str = ComputePoolNameArgument,
+    name: FQN = ComputePoolNameArgument,
     min_nodes: Optional[int] = MinNodesOption(default=None, show_default=False),
     max_nodes: Optional[int] = MaxNodesOption(show_default=False),
     auto_resume: Optional[bool] = AutoResumeOption(default=None, show_default=False),
@@ -187,7 +188,7 @@ def set_property(
     Sets one or more properties for the compute pool.
     """
     cursor = ComputePoolManager().set_property(
-        pool_name=name,
+        pool_name=name.identifier,
         min_nodes=min_nodes,
         max_nodes=max_nodes,
         auto_resume=auto_resume,
@@ -199,7 +200,7 @@ def set_property(
 
 @app.command("unset", requires_connection=True)
 def unset_property(
-    name: str = ComputePoolNameArgument,
+    name: FQN = ComputePoolNameArgument,
     auto_resume: bool = AutoResumeOption(
         default=False,
         param_decls=["--auto-resume"],
@@ -223,7 +224,7 @@ def unset_property(
     Resets one or more properties for the compute pool to their default value(s).
     """
     cursor = ComputePoolManager().unset_property(
-        pool_name=name,
+        pool_name=name.identifier,
         auto_resume=auto_resume,
         auto_suspend_secs=auto_suspend_secs,
         comment=comment,
@@ -232,9 +233,9 @@ def unset_property(
 
 
 @app.command(requires_connection=True)
-def status(pool_name: str = ComputePoolNameArgument, **options) -> CommandResult:
+def status(pool_name: FQN = ComputePoolNameArgument, **options) -> CommandResult:
     """
     Retrieves the status of a compute pool along with a relevant message, if one exists.
     """
-    cursor = ComputePoolManager().status(pool_name=pool_name)
+    cursor = ComputePoolManager().status(pool_name=pool_name.identifier)
     return SingleQueryResult(cursor)

snowflake/cli/{plugins → _plugins}/spcs/compute_pool/manager.py

@@ -16,13 +16,13 @@ from __future__ import annotations
 
 from typing import List, Optional
 
-from snowflake.cli.api.constants import ObjectType
-from snowflake.cli.api.sql_execution import SqlExecutionMixin
-from snowflake.cli.plugins.spcs.common import (
+from snowflake.cli._plugins.spcs.common import (
     NoPropertiesProvidedError,
     handle_object_already_exists,
     strip_empty_lines,
 )
+from snowflake.cli.api.constants import ObjectType
+from snowflake.cli.api.sql_execution import SqlExecutionMixin
 from snowflake.connector.cursor import SnowflakeCursor
 from snowflake.connector.errors import ProgrammingError
 

snowflake/cli/{plugins → _plugins}/spcs/image_registry/commands.py

@@ -12,11 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
-from snowflake.cli.api.output.types import MessageResult, ObjectResult
-from snowflake.cli.plugins.spcs.image_registry.manager import (
+from snowflake.cli._plugins.spcs.image_registry.manager import (
     RegistryManager,
 )
+from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
+from snowflake.cli.api.output.types import MessageResult, ObjectResult
 
 app = SnowTyperFactory(
     name="image-registry",

snowflake/cli/{plugins → _plugins}/spcs/image_repository/commands.py

@@ -20,26 +20,28 @@ from typing import Optional
 import requests
 import typer
 from click import ClickException
+from snowflake.cli._plugins.object.command_aliases import (
+    add_object_command_aliases,
+    scope_option,
+)
+from snowflake.cli._plugins.spcs.image_registry.manager import RegistryManager
+from snowflake.cli._plugins.spcs.image_repository.manager import ImageRepositoryManager
 from snowflake.cli.api.commands.flags import (
     IfNotExistsOption,
     ReplaceOption,
+    identifier_argument,
     like_option,
 )
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
 from snowflake.cli.api.console import cli_console
 from snowflake.cli.api.constants import ObjectType
+from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.output.types import (
     CollectionResult,
     MessageResult,
     SingleQueryResult,
 )
 from snowflake.cli.api.project.util import is_valid_object_name
-from snowflake.cli.plugins.object.command_aliases import (
-    add_object_command_aliases,
-    scope_option,
-)
-from snowflake.cli.plugins.spcs.image_registry.manager import RegistryManager
-from snowflake.cli.plugins.spcs.image_repository.manager import ImageRepositoryManager
 
 app = SnowTyperFactory(
     name="image-repository",
@@ -48,18 +50,18 @@ app = SnowTyperFactory(
 )
 
 
-def _repo_name_callback(name: str):
-    if not is_valid_object_name(name, max_depth=2, allow_quoted=False):
+def _repo_name_callback(name: FQN):
+    if not is_valid_object_name(name.identifier, max_depth=2, allow_quoted=False):
         raise ClickException(
             f"'{name}' is not a valid image repository name. Note that image repository names must be unquoted identifiers. The same constraint also applies to database and schema names where you create an image repository."
         )
     return name
 
 
-REPO_NAME_ARGUMENT = typer.Argument(
-    help="Name of the image repository.",
+REPO_NAME_ARGUMENT = identifier_argument(
+    sf_object="image repository",
+    example="my_repository",
     callback=_repo_name_callback,
-    show_default=False,
 )
 
 add_object_command_aliases(
@@ -76,7 +78,7 @@ add_object_command_aliases(
 
 @app.command(requires_connection=True)
 def create(
-    name: str = REPO_NAME_ARGUMENT,
+    name: FQN = REPO_NAME_ARGUMENT,
     replace: bool = ReplaceOption(),
     if_not_exists: bool = IfNotExistsOption(),
     **options,
@@ -86,21 +88,21 @@ def create(
     """
     return SingleQueryResult(
         ImageRepositoryManager().create(
-            name=name, replace=replace, if_not_exists=if_not_exists
+            name=name.identifier, replace=replace, if_not_exists=if_not_exists
         )
     )
 
 
 @app.command("list-images", requires_connection=True)
 def list_images(
-    name: str = REPO_NAME_ARGUMENT,
+    name: FQN = REPO_NAME_ARGUMENT,
     **options,
 ) -> CollectionResult:
     """Lists images in the given repository."""
     repository_manager = ImageRepositoryManager()
     database = repository_manager.get_database()
     schema = repository_manager.get_schema()
-    url = repository_manager.get_repository_url(name)
+    url = repository_manager.get_repository_url(name.identifier)
     api_url = repository_manager.get_repository_api_url(url)
     bearer_login = RegistryManager().login_to_registry(api_url)
     repos = []
@@ -136,7 +138,7 @@ def list_images(
 
 @app.command("list-tags", requires_connection=True)
 def list_tags(
-    name: str = REPO_NAME_ARGUMENT,
+    name: FQN = REPO_NAME_ARGUMENT,
     image_name: str = typer.Option(
         ...,
         "--image-name",
@@ -150,7 +152,7 @@ def list_tags(
     """Lists tags for the given image in a repository."""
 
     repository_manager = ImageRepositoryManager()
-    url = repository_manager.get_repository_url(name)
+    url = repository_manager.get_repository_url(name.identifier)
     api_url = repository_manager.get_repository_api_url(url)
     bearer_login = RegistryManager().login_to_registry(api_url)
 
@@ -187,10 +189,14 @@ def list_tags(
 
 @app.command("url", requires_connection=True)
 def repo_url(
-    name: str = REPO_NAME_ARGUMENT,
+    name: FQN = REPO_NAME_ARGUMENT,
     **options,
 ):
     """Returns the URL for the given repository."""
     return MessageResult(
-        (ImageRepositoryManager().get_repository_url(repo_name=name, with_scheme=False))
+        (
+            ImageRepositoryManager().get_repository_url(
+                repo_name=name.identifier, with_scheme=False
+            )
+        )
     )

snowflake/cli/{plugins → _plugins}/spcs/image_repository/manager.py

@@ -14,10 +14,10 @@
 
 from urllib.parse import urlparse
 
+from snowflake.cli._plugins.spcs.common import handle_object_already_exists
 from snowflake.cli.api.constants import ObjectType
 from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.sql_execution import SqlExecutionMixin
-from snowflake.cli.plugins.spcs.common import handle_object_already_exists
 from snowflake.connector.errors import ProgrammingError
 
 

snowflake/cli/{plugins → _plugins}/spcs/plugin_spec.py

@@ -12,13 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from snowflake.cli._plugins.spcs import app
 from snowflake.cli.api.plugins.command import (
     SNOWCLI_ROOT_COMMAND_PATH,
     CommandSpec,
     CommandType,
     plugin_hook_impl,
 )
-from snowflake.cli.plugins.spcs import app
 
 
 @plugin_hook_impl

snowflake/cli/{plugins → _plugins}/spcs/services/commands.py

@@ -20,13 +20,25 @@ from typing import List, Optional
 
 import typer
 from click import ClickException
+from snowflake.cli._plugins.object.command_aliases import (
+    add_object_command_aliases,
+    scope_option,
+)
+from snowflake.cli._plugins.object.common import CommentOption, Tag, TagOption
+from snowflake.cli._plugins.spcs.common import (
+    print_log_lines,
+    validate_and_set_instances,
+)
+from snowflake.cli._plugins.spcs.services.manager import ServiceManager
 from snowflake.cli.api.commands.flags import (
     IfNotExistsOption,
     OverrideableOption,
+    identifier_argument,
     like_option,
 )
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
 from snowflake.cli.api.constants import ObjectType
+from snowflake.cli.api.identifiers import FQN
 from snowflake.cli.api.output.types import (
     CommandResult,
     QueryJsonValueResult,
@@ -34,16 +46,6 @@ from snowflake.cli.api.output.types import (
     SingleQueryResult,
 )
 from snowflake.cli.api.project.util import is_valid_object_name
-from snowflake.cli.plugins.object.command_aliases import (
-    add_object_command_aliases,
-    scope_option,
-)
-from snowflake.cli.plugins.object.common import CommentOption, Tag, TagOption
-from snowflake.cli.plugins.spcs.common import (
-    print_log_lines,
-    validate_and_set_instances,
-)
-from snowflake.cli.plugins.spcs.services.manager import ServiceManager
 
 app = SnowTyperFactory(
     name="service",
@@ -52,19 +54,18 @@ app = SnowTyperFactory(
 )
 
 
-def _service_name_callback(name: str) -> str:
-    if not is_valid_object_name(name, max_depth=2, allow_quoted=False):
+def _service_name_callback(name: FQN) -> FQN:
+    if not is_valid_object_name(name.identifier, max_depth=2, allow_quoted=False):
         raise ClickException(
             f"'{name}' is not a valid service name. Note service names must be unquoted identifiers. The same constraint also applies to database and schema names where you create a service."
         )
     return name
 
 
-ServiceNameArgument = typer.Argument(
-    ...,
-    help="Name of the service.",
+ServiceNameArgument = identifier_argument(
+    sf_object="service",
+    example="my_service",
     callback=_service_name_callback,
-    show_default=False,
 )
 
 SpecPathOption = typer.Option(
@@ -116,7 +117,7 @@ add_object_command_aliases(
 
 @app.command(requires_connection=True)
 def create(
-    name: str = ServiceNameArgument,
+    name: FQN = ServiceNameArgument,
     compute_pool: str = typer.Option(
         ...,
         "--compute-pool",
@@ -145,7 +146,7 @@ def create(
         min_instances, max_instances, "instances"
     )
     cursor = ServiceManager().create(
-        service_name=name,
+        service_name=name.identifier,
         min_instances=min_instances,
         max_instances=max_instances,
         compute_pool=compute_pool,
@@ -161,17 +162,50 @@ def create(
 
 
 @app.command(requires_connection=True)
-def status(name: str = ServiceNameArgument, **options) -> CommandResult:
+def execute_job(
+    name: FQN = ServiceNameArgument,
+    compute_pool: str = typer.Option(
+        ...,
+        "--compute-pool",
+        help="Compute pool to run the job service on.",
+        show_default=False,
+    ),
+    spec_path: Path = SpecPathOption,
+    external_access_integrations: Optional[List[str]] = typer.Option(
+        None,
+        "--eai-name",
+        help="Identifies External Access Integrations(EAI) that the job service can access. This option may be specified multiple times for multiple EAIs.",
+    ),
+    query_warehouse: Optional[str] = QueryWarehouseOption(),
+    comment: Optional[str] = CommentOption(help=_COMMENT_HELP),
+    **options,
+) -> CommandResult:
+    """
+    Creates and executes a job service in the current schema.
+    """
+    cursor = ServiceManager().execute_job(
+        job_service_name=name.identifier,
+        compute_pool=compute_pool,
+        spec_path=spec_path,
+        external_access_integrations=external_access_integrations,
+        query_warehouse=query_warehouse,
+        comment=comment,
+    )
+    return SingleQueryResult(cursor)
+
+
+@app.command(requires_connection=True)
+def status(name: FQN = ServiceNameArgument, **options) -> CommandResult:
     """
     Retrieves the status of a service.
     """
-    cursor = ServiceManager().status(service_name=name)
+    cursor = ServiceManager().status(service_name=name.identifier)
     return QueryJsonValueResult(cursor)
 
 
 @app.command(requires_connection=True)
 def logs(
-    name: str = ServiceNameArgument,
+    name: FQN = ServiceNameArgument,
     container_name: str = typer.Option(
         ...,
         "--container-name",
@@ -193,7 +227,7 @@ def logs(
     Retrieves local logs from a service container.
     """
     results = ServiceManager().logs(
-        service_name=name,
+        service_name=name.identifier,
         instance_id=instance_id,
         container_name=container_name,
         num_lines=num_lines,
@@ -205,7 +239,7 @@ def logs(
 
 @app.command(requires_connection=True)
 def upgrade(
-    name: str = ServiceNameArgument,
+    name: FQN = ServiceNameArgument,
     spec_path: Path = SpecPathOption,
     **options,
 ):
@@ -213,20 +247,20 @@ def upgrade(
     Updates an existing service with a new specification file.
     """
     return SingleQueryResult(
-        ServiceManager().upgrade_spec(service_name=name, spec_path=spec_path)
+        ServiceManager().upgrade_spec(service_name=name.identifier, spec_path=spec_path)
     )
 
 
 @app.command("list-endpoints", requires_connection=True)
-def list_endpoints(name: str = ServiceNameArgument, **options):
+def list_endpoints(name: FQN = ServiceNameArgument, **options):
     """
     Lists the endpoints in a service.
     """
-    return QueryResult(ServiceManager().list_endpoints(service_name=name))
+    return QueryResult(ServiceManager().list_endpoints(service_name=name.identifier))
 
 
 @app.command(requires_connection=True)
-def suspend(name: str = ServiceNameArgument, **options) -> CommandResult:
+def suspend(name: FQN = ServiceNameArgument, **options) -> CommandResult:
     """
     Suspends the service, shutting down and deleting all its containers.
     """
@@ -234,7 +268,7 @@ def suspend(name: str = ServiceNameArgument, **options) -> CommandResult:
 
 
 @app.command(requires_connection=True)
-def resume(name: str = ServiceNameArgument, **options) -> CommandResult:
+def resume(name: FQN = ServiceNameArgument, **options) -> CommandResult:
     """
     Resumes the service from a SUSPENDED state.
     """
@@ -243,7 +277,7 @@ def resume(name: str = ServiceNameArgument, **options) -> CommandResult:
 
 @app.command("set", requires_connection=True)
 def set_property(
-    name: str = ServiceNameArgument,
+    name: FQN = ServiceNameArgument,
     min_instances: Optional[int] = MinInstancesOption(default=None, show_default=False),
     max_instances: Optional[int] = MaxInstancesOption(show_default=False),
     query_warehouse: Optional[str] = QueryWarehouseOption(show_default=False),
@@ -255,7 +289,7 @@ def set_property(
     Sets one or more properties for the service.
     """
     cursor = ServiceManager().set_property(
-        service_name=name,
+        service_name=name.identifier,
         min_instances=min_instances,
         max_instances=max_instances,
         query_warehouse=query_warehouse,
@@ -267,7 +301,7 @@ def set_property(
 
 @app.command("unset", requires_connection=True)
 def unset_property(
-    name: str = ServiceNameArgument,
+    name: FQN = ServiceNameArgument,
     min_instances: bool = MinInstancesOption(
         default=False,
         help=f"Reset the MIN_INSTANCES property - {_MIN_INSTANCES_HELP}",
@@ -301,7 +335,7 @@ def unset_property(
     Resets one or more properties for the service to their default value(s).
     """
     cursor = ServiceManager().unset_property(
-        service_name=name,
+        service_name=name.identifier,
         min_instances=min_instances,
         max_instances=max_instances,
         query_warehouse=query_warehouse,