snowflake-cli-labs 2.8.0rc1__py3-none-any.whl → 3.0.0__py3-none-any.whl

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,

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

@@ -19,15 +19,15 @@ from pathlib import Path
 from typing import List, Optional
 
 import yaml
-from snowflake.cli.api.constants import DEFAULT_SIZE_LIMIT_MB, ObjectType
-from snowflake.cli.api.secure_path import SecurePath
-from snowflake.cli.api.sql_execution import SqlExecutionMixin
-from snowflake.cli.plugins.object.common import Tag
-from snowflake.cli.plugins.spcs.common import (
+from snowflake.cli._plugins.object.common import Tag
+from snowflake.cli._plugins.spcs.common import (
     NoPropertiesProvidedError,
     handle_object_already_exists,
     strip_empty_lines,
 )
+from snowflake.cli.api.constants import DEFAULT_SIZE_LIMIT_MB, ObjectType
+from snowflake.cli.api.secure_path import SecurePath
+from snowflake.cli.api.sql_execution import SqlExecutionMixin
 from snowflake.connector.cursor import SnowflakeCursor
 from snowflake.connector.errors import ProgrammingError
 
@@ -85,6 +85,44 @@ class ServiceManager(SqlExecutionMixin):
         except ProgrammingError as e:
             handle_object_already_exists(e, ObjectType.SERVICE, service_name)
 
+    def execute_job(
+        self,
+        job_service_name: str,
+        compute_pool: str,
+        spec_path: Path,
+        external_access_integrations: Optional[List[str]],
+        query_warehouse: Optional[str],
+        comment: Optional[str],
+    ) -> SnowflakeCursor:
+        spec = self._read_yaml(spec_path)
+        query = f"""\
+                EXECUTE JOB SERVICE
+                IN COMPUTE POOL {compute_pool}
+                FROM SPECIFICATION $$
+                {spec}
+                $$
+                NAME = {job_service_name}
+                """.splitlines()
+
+        if external_access_integrations:
+            external_access_integration_list = ",".join(
+                f"{e}" for e in external_access_integrations
+            )
+            query.append(
+                f"EXTERNAL_ACCESS_INTEGRATIONS = ({external_access_integration_list})"
+            )
+
+        if query_warehouse:
+            query.append(f"QUERY_WAREHOUSE = {query_warehouse}")
+
+        if comment:
+            query.append(f"COMMENT = {comment}")
+
+        try:
+            return self._execute_query(strip_empty_lines(query))
+        except ProgrammingError as e:
+            handle_object_already_exists(e, ObjectType.SERVICE, job_service_name)
+
     def _read_yaml(self, path: Path) -> str:
         # TODO(aivanou): Add validation towards schema
         with SecurePath(path).open("r", read_file_limit_mb=DEFAULT_SIZE_LIMIT_MB) as fh:

snowflake/cli/{plugins → _plugins}/sql/commands.py

@@ -17,43 +17,46 @@ from __future__ import annotations
 from pathlib import Path
 from typing import List, Optional
 
-import typer
+from snowflake.cli._plugins.sql.manager import SqlManager
 from snowflake.cli.api.commands.decorators import with_project_definition
 from snowflake.cli.api.commands.flags import (
-    parse_key_value_variables,
     variables_option,
 )
+from snowflake.cli.api.commands.overrideable_parameter import OverrideableOption
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
+from snowflake.cli.api.commands.utils import parse_key_value_variables
 from snowflake.cli.api.output.types import CommandResult, MultipleResults, QueryResult
-from snowflake.cli.plugins.sql.manager import SqlManager
 
 # simple Typer with defaults because it won't become a command group as it contains only one command
 app = SnowTyperFactory()
 
+SOURCE_EXCLUSIVE_OPTIONS_NAMES = ["query", "files", "std_in"]
+
+SourceOption = OverrideableOption(
+    mutually_exclusive=SOURCE_EXCLUSIVE_OPTIONS_NAMES, show_default=False
+)
+
 
 @app.command(name="sql", requires_connection=True, no_args_is_help=True)
 @with_project_definition(is_optional=True)
 def execute_sql(
-    query: Optional[str] = typer.Option(
-        None,
-        "--query",
-        "-q",
+    query: Optional[str] = SourceOption(
+        default=None,
+        param_decls=["--query", "-q"],
         help="Query to execute.",
     ),
-    files: Optional[List[Path]] = typer.Option(
-        None,
-        "--filename",
-        "-f",
+    files: Optional[List[Path]] = SourceOption(
+        default=[],
+        param_decls=["--filename", "-f"],
         exists=True,
         file_okay=True,
         dir_okay=False,
         readable=True,
         help="File to execute.",
     ),
-    std_in: Optional[bool] = typer.Option(
-        False,
-        "--stdin",
-        "-i",
+    std_in: Optional[bool] = SourceOption(
+        default=False,
+        param_decls=["--stdin", "-i"],
         help="Read the query from standard input. Use it when piping input to this command.",
     ),
     data_override: List[str] = variables_option(
@@ -70,9 +73,9 @@ def execute_sql(
     Query to execute can be specified using query option, filename option (all queries from file will be executed)
     or via stdin by piping output from other command. For example `cat my.sql | snow sql -i`.
 
-    The command supports variable substitution that happens on client-side. Both &VARIABLE or &{ VARIABLE }
-    syntax are supported.
+    The command supports variable substitution that happens on client-side.
     """
+
     data = {}
     if data_override:
         data = {v.key: v.value for v in parse_key_value_variables(data_override)}

snowflake/cli/{plugins → _plugins}/sql/manager.py

@@ -22,10 +22,10 @@ from typing import Dict, Iterable, List, Tuple
 
 from click import ClickException, UsageError
 from jinja2 import UndefinedError
+from snowflake.cli._plugins.sql.snowsql_templating import transpile_snowsql_templates
 from snowflake.cli.api.rendering.sql_templates import snowflake_sql_jinja_render
 from snowflake.cli.api.secure_path import UNLIMITED, SecurePath
 from snowflake.cli.api.sql_execution import SqlExecutionMixin, VerboseCursor
-from snowflake.cli.plugins.sql.snowsql_templating import transpile_snowsql_templates
 from snowflake.connector.cursor import SnowflakeCursor
 from snowflake.connector.util_text import split_statements
 

snowflake/cli/{plugins → _plugins}/sql/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.sql import commands
 from snowflake.cli.api.plugins.command import (
     SNOWCLI_ROOT_COMMAND_PATH,
     CommandSpec,
     CommandType,
     plugin_hook_impl,
 )
-from snowflake.cli.plugins.sql import commands
 
 
 @plugin_hook_impl

snowflake/cli/{plugins → _plugins}/stage/commands.py

@@ -21,16 +21,29 @@ from typing import List, Optional
 
 import click
 import typer
-from snowflake.cli.api.cli_global_context import cli_context
+from snowflake.cli._plugins.object.command_aliases import (
+    add_object_command_aliases,
+    scope_option,
+)
+from snowflake.cli._plugins.stage.diff import (
+    DiffResult,
+    compute_stage_diff,
+)
+from snowflake.cli._plugins.stage.manager import StageManager
+from snowflake.cli._plugins.stage.utils import print_diff_to_console
+from snowflake.cli.api.cli_global_context import get_cli_context
+from snowflake.cli.api.commands.common import OnErrorType
 from snowflake.cli.api.commands.flags import (
     ExecuteVariablesOption,
     OnErrorOption,
     PatternOption,
+    identifier_stage_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.formats import OutputFormat
 from snowflake.cli.api.output.types import (
     CollectionResult,
@@ -40,23 +53,13 @@ from snowflake.cli.api.output.types import (
     SingleQueryResult,
 )
 from snowflake.cli.api.utils.path_utils import is_stage_path
-from snowflake.cli.plugins.object.command_aliases import (
-    add_object_command_aliases,
-    scope_option,
-)
-from snowflake.cli.plugins.stage.diff import (
-    DiffResult,
-    compute_stage_diff,
-    print_diff_to_console,
-)
-from snowflake.cli.plugins.stage.manager import OnErrorType, StageManager
 
 app = SnowTyperFactory(
     name="stage",
     help="Manages stages.",
 )
 
-StageNameArgument = typer.Argument(..., help="Name of the stage.", show_default=False)
+StageNameArgument = identifier_stage_argument(sf_object="stage", example="@my_stage")
 
 add_object_command_aliases(
     app=app,
@@ -142,17 +145,17 @@ def copy(
 
 
 @app.command("create", requires_connection=True)
-def stage_create(stage_name: str = StageNameArgument, **options) -> CommandResult:
+def stage_create(stage_name: FQN = StageNameArgument, **options) -> CommandResult:
     """
     Creates a named stage if it does not already exist.
     """
-    cursor = StageManager().create(stage_name=stage_name)
+    cursor = StageManager().create(fqn=stage_name)
     return SingleQueryResult(cursor)
 
 
 @app.command("remove", requires_connection=True)
 def stage_remove(
-    stage_name: str = StageNameArgument,
+    stage_name: FQN = StageNameArgument,
     file_name: str = typer.Argument(
         ...,
         help="Name of the file to remove.",
@@ -164,7 +167,7 @@ def stage_remove(
     Removes a file from a stage.
     """
 
-    cursor = StageManager().remove(stage_name=stage_name, path=file_name)
+    cursor = StageManager().remove(stage_name=stage_name.identifier, path=file_name)
     return SingleQueryResult(cursor)
 
 
@@ -186,7 +189,7 @@ def stage_diff(
     diff: DiffResult = compute_stage_diff(
         local_root=Path(folder_name), stage_fqn=stage_name
     )
-    if cli_context.output_format == OutputFormat.JSON:
+    if get_cli_context().output_format == OutputFormat.JSON:
         return ObjectResult(diff.to_dict())
     else:
         print_diff_to_console(diff)