snowflake-cli-labs 2.6.0rc0__py3-none-any.whl → 2.7.0rc0__py3-none-any.whl

snowflake/cli/api/rendering/sql_templates.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2024 Snowflake Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import Dict, Optional
+
+from click import ClickException
+from jinja2 import StrictUndefined, loaders
+from snowflake.cli.api.cli_global_context import cli_context
+from snowflake.cli.api.rendering.jinja import (
+    CONTEXT_KEY,
+    IgnoreAttrEnvironment,
+    env_bootstrap,
+)
+
+_SQL_TEMPLATE_START = "&{"
+_SQL_TEMPLATE_END = "}"
+
+
+def get_sql_cli_jinja_env(*, loader: Optional[loaders.BaseLoader] = None):
+    _random_block = "___very___unique___block___to___disable___logic___blocks___"
+    return env_bootstrap(
+        IgnoreAttrEnvironment(
+            loader=loader or loaders.BaseLoader(),
+            keep_trailing_newline=True,
+            variable_start_string=_SQL_TEMPLATE_START,
+            variable_end_string=_SQL_TEMPLATE_END,
+            block_start_string=_random_block,
+            block_end_string=_random_block,
+            undefined=StrictUndefined,
+        )
+    )
+
+
+def snowflake_sql_jinja_render(content: str, data: Dict | None = None) -> str:
+    data = data or {}
+    if CONTEXT_KEY in data:
+        raise ClickException(
+            f"{CONTEXT_KEY} in user defined data. The `{CONTEXT_KEY}` variable is reserved for CLI usage."
+        )
+
+    context_data = cli_context.template_context
+    context_data.update(data)
+    return get_sql_cli_jinja_env().from_string(content).render(**context_data)

snowflake/cli/api/rest_api.py

@@ -16,16 +16,27 @@ from __future__ import annotations
 
 import json
 import logging
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, Optional
 
+from click import ClickException
 from snowflake.cli.api.constants import SF_REST_API_URL_PREFIX
 from snowflake.connector.connection import SnowflakeConnection
-from snowflake.connector.errors import InterfaceError
+from snowflake.connector.errors import BadRequest, InterfaceError
 from snowflake.connector.network import SnowflakeRestful
 
 log = logging.getLogger(__name__)
 
 
+def _pluralize_object_type(object_type: str) -> str:
+    """
+    Pluralize object type without depending on OBJECT_TO_NAMES.
+    """
+    if object_type.endswith("y"):
+        return object_type[:-1].lower() + "ies"
+    else:
+        return object_type.lower() + "s"
+
+
 class RestApi:
     def __init__(self, connection: SnowflakeConnection):
         self.conn = connection
@@ -43,6 +54,13 @@ class RestApi:
                 return False
             raise err
 
+    def _fetch_endpoint_exists(self, url: str) -> bool:
+        try:
+            result = self.send_rest_request(url, method="get")
+            return bool(result)
+        except BadRequest:
+            return False
+
     def send_rest_request(
         self, url: str, method: str, data: Optional[Dict[str, Any]] = None
     ):
@@ -75,12 +93,18 @@ class RestApi:
             no_retry=True,
         )
 
-    def determine_url_for_create_query(
-        self, *, plural_object_type: str
-    ) -> Optional[str]:
+    def _database_exists(self, db_name: str) -> bool:
+        url = f"{SF_REST_API_URL_PREFIX}/databases/{db_name}"
+        return self._fetch_endpoint_exists(url)
+
+    def _schema_exists(self, db_name: str, schema_name: str) -> bool:
+        url = f"{SF_REST_API_URL_PREFIX}/databases/{db_name}/schemas/{schema_name}"
+        return self._fetch_endpoint_exists(url)
+
+    def determine_url_for_create_query(self, object_type: str) -> str:
         """
         Determine an url for creating an object of given type via REST API.
-        The function returns None if URL cannot be determined.
+        If URL cannot be determined, the function throws CannotDetermineCreateURLException exception.
 
         URLs we check:
          * /api/v2/<type>/
@@ -92,22 +116,57 @@ class RestApi:
          To check whether an URL exists, we send read-only GET request (LIST endpoint,
          which should imply CREATE endpoint).
         """
-        urls_to_be_checked: List[Optional[str]] = [
-            f"{SF_REST_API_URL_PREFIX}/{plural_object_type}/",
-            (
-                f"{SF_REST_API_URL_PREFIX}/databases/{self.conn.database}/{plural_object_type}/"
-                if self.conn.database
-                else None
-            ),
-            (
-                f"{SF_REST_API_URL_PREFIX}/databases/{self.conn.database}/schemas/{self.conn.schema}/{plural_object_type}/"
-                if self.conn.database and self.conn.schema
-                else None
-            ),
-        ]
-
-        for url in urls_to_be_checked:
-            if url and self.get_endpoint_exists(url):
-                return url
-
-        return None
+        plural_object_type = _pluralize_object_type(object_type)
+
+        if self.get_endpoint_exists(
+            url := f"{SF_REST_API_URL_PREFIX}/{plural_object_type}/"
+        ):
+            return url
+
+        db = self.conn.database
+        if not db:
+            raise DatabaseNotDefinedException(
+                "Database not defined in connection. Please try again with `--database` flag."
+            )
+        if not self._database_exists(db):
+            raise DatabaseNotExistsException(f"Database '{db}' does not exist.")
+        if self.get_endpoint_exists(
+            url := f"{SF_REST_API_URL_PREFIX}/databases/{db}/{plural_object_type}/"
+        ):
+            return url
+
+        schema = self.conn.schema
+        if not schema:
+            raise SchemaNotDefinedException(
+                "Schema not defined in connection. Please try again with `--schema` flag."
+            )
+        if not self._schema_exists(db_name=db, schema_name=schema):
+            raise SchemaNotExistsException(f"Schema '{schema}' does not exist.")
+        if self.get_endpoint_exists(
+            url := f"{SF_REST_API_URL_PREFIX}/databases/{self.conn.database}/schemas/{self.conn.schema}/{plural_object_type}/"
+        ):
+            return url
+
+        raise CannotDetermineCreateURLException(
+            f"Create operation for type {object_type} is not supported. Try using `sql -q 'CREATE ...'` command."
+        )
+
+
+class DatabaseNotDefinedException(ClickException):
+    pass
+
+
+class SchemaNotDefinedException(ClickException):
+    pass
+
+
+class DatabaseNotExistsException(ClickException):
+    pass
+
+
+class SchemaNotExistsException(ClickException):
+    pass
+
+
+class CannotDetermineCreateURLException(ClickException):
+    pass

snowflake/cli/api/sql_execution.py

@@ -41,12 +41,22 @@ from snowflake.connector.errors import ProgrammingError
 
 class SqlExecutionMixin:
     def __init__(self):
-        pass
+        self._snowpark_session = None
 
     @property
     def _conn(self):
         return cli_context.connection
 
+    @property
+    def snowpark_session(self):
+        if not self._snowpark_session:
+            from snowflake.snowpark.session import Session
+
+            self._snowpark_session = Session.builder.configs(
+                {"connection": self._conn}
+            ).create()
+        return self._snowpark_session
+
     @cached_property
     def _log(self):
         return logging.getLogger(__name__)
@@ -107,6 +117,35 @@ class SqlExecutionMixin:
             if is_different_role:
                 self._execute_query(f"use role {prev_role}")
 
+    @contextmanager
+    def use_warehouse(self, new_wh: str):
+        """
+        Switches to a different warehouse for a while, then switches back.
+        This is a no-op if the requested warehouse is already active.
+        If there is no default warehouse in the account, it will throw an error.
+        """
+
+        wh_result = self._execute_query(
+            f"select current_warehouse()", cursor_class=DictCursor
+        ).fetchone()
+        # If user has an assigned default warehouse, prev_wh will contain a value even if the warehouse is suspended.
+        try:
+            prev_wh = wh_result["CURRENT_WAREHOUSE()"]
+        except:
+            prev_wh = None
+
+        # new_wh is not None, and should already be a valid identifier, no additional check is performed here.
+        is_different_wh = new_wh != prev_wh
+        try:
+            if is_different_wh:
+                self._log.debug("Using warehouse: %s", new_wh)
+                self.use(object_type=ObjectType.WAREHOUSE, name=new_wh)
+            yield
+        finally:
+            if prev_wh and is_different_wh:
+                self._log.debug("Switching back to warehouse: %s", prev_wh)
+                self.use(object_type=ObjectType.WAREHOUSE, name=prev_wh)
+
     def create_password_secret(
         self, name: str, username: str, password: str
     ) -> SnowflakeCursor:

snowflake/cli/api/utils/definition_rendering.py

@@ -22,13 +22,16 @@ from packaging.version import Version
 from snowflake.cli.api.console import cli_console as cc
 from snowflake.cli.api.exceptions import CycleDetectedError, InvalidTemplate
 from snowflake.cli.api.project.schemas.project_definition import (
-    ProjectDefinition,
     ProjectProperties,
+    build_project_definition,
+)
+from snowflake.cli.api.rendering.jinja import CONTEXT_KEY
+from snowflake.cli.api.rendering.project_definition_templates import (
+    get_project_definition_cli_jinja_env,
 )
 from snowflake.cli.api.utils.dict_utils import traverse
 from snowflake.cli.api.utils.graph import Graph, Node
 from snowflake.cli.api.utils.models import ProjectEnvironment
-from snowflake.cli.api.utils.rendering import CONTEXT_KEY, get_snowflake_cli_jinja_env
 from snowflake.cli.api.utils.types import Context, Definition
 
 
@@ -286,7 +289,7 @@ def render_definition_template(
         return ProjectProperties(None, {CONTEXT_KEY: {"env": environment_overrides}})
 
     project_context = {CONTEXT_KEY: definition}
-    template_env = TemplatedEnvironment(get_snowflake_cli_jinja_env())
+    template_env = TemplatedEnvironment(get_project_definition_cli_jinja_env())
 
     if "definition_version" not in definition or Version(
         definition["definition_version"]
@@ -301,7 +304,7 @@ def render_definition_template(
             # also warn on Exception, as it means the user is incorrectly attempting to use templating
             _template_version_warning()
 
-        project_definition = ProjectDefinition(**original_definition)
+        project_definition = build_project_definition(**original_definition)
         project_context[CONTEXT_KEY]["env"] = environment_overrides
         return ProjectProperties(project_definition, project_context)
 
@@ -337,5 +340,5 @@ def render_definition_template(
 
     definition["env"] = ProjectEnvironment(default_env, override_env)
     project_context[CONTEXT_KEY] = definition
-    project_definition = ProjectDefinition(**definition)
+    project_definition = build_project_definition(**definition)
     return ProjectProperties(project_definition, project_context)

snowflake/cli/app/cli_app.py

@@ -30,7 +30,6 @@ from snowflake.cli.api.config import config_init
 from snowflake.cli.api.output.formats import OutputFormat
 from snowflake.cli.api.output.types import CollectionResult
 from snowflake.cli.api.secure_path import SecurePath
-from snowflake.cli.app import loggers
 from snowflake.cli.app.api_impl.plugin.plugin_config_provider_impl import (
     PluginConfigProviderImpl,
 )
@@ -89,7 +88,6 @@ def _commands_registration_callback(value: bool):
 @_commands_registration.before
 def _config_init_callback(configuration_file: Optional[Path]):
     config_init(configuration_file)
-    loggers.create_initial_loggers()
 
 
 @_commands_registration.before

snowflake/cli/app/commands_registration/builtin_plugins.py

@@ -15,6 +15,7 @@
 from snowflake.cli.plugins.connection import plugin_spec as connection_plugin_spec
 from snowflake.cli.plugins.cortex import plugin_spec as cortex_plugin_spec
 from snowflake.cli.plugins.git import plugin_spec as git_plugin_spec
+from snowflake.cli.plugins.init import plugin_spec as init_plugin_spec
 from snowflake.cli.plugins.nativeapp import plugin_spec as nativeapp_plugin_spec
 from snowflake.cli.plugins.notebook import plugin_spec as notebook_plugin_spec
 from snowflake.cli.plugins.object import plugin_spec as object_plugin_spec
@@ -28,6 +29,7 @@ from snowflake.cli.plugins.spcs import plugin_spec as spcs_plugin_spec
 from snowflake.cli.plugins.sql import plugin_spec as sql_plugin_spec
 from snowflake.cli.plugins.stage import plugin_spec as stage_plugin_spec
 from snowflake.cli.plugins.streamlit import plugin_spec as streamlit_plugin_spec
+from snowflake.cli.plugins.workspace import plugin_spec as workspace_plugin_spec
 
 
 # plugin name to plugin spec
@@ -45,6 +47,8 @@ def get_builtin_plugin_name_to_plugin_spec():
         "notebook": notebook_plugin_spec,
         "object-stage-deprecated": object_stage_deprecated_plugin_spec,
         "cortex": cortex_plugin_spec,
+        "init": init_plugin_spec,
+        "workspace": workspace_plugin_spec,
     }
 
     return plugin_specs

snowflake/cli/app/dev/docs/project_definition_docs_generator.py

@@ -18,7 +18,7 @@ import logging
 from typing import Any, Dict
 
 from pydantic.json_schema import model_json_schema
-from snowflake.cli.api.project.schemas.project_definition import ProjectDefinition
+from snowflake.cli.api.project.schemas.project_definition import DefinitionV11
 from snowflake.cli.api.secure_path import SecurePath
 from snowflake.cli.app.dev.docs.project_definition_generate_json_schema import (
     ProjectDefinitionGenerateJsonSchema,
@@ -39,7 +39,7 @@ def generate_project_definition_docs(root: SecurePath):
 
     root.mkdir(exist_ok=True)
     list_of_sections = model_json_schema(
-        ProjectDefinition, schema_generator=ProjectDefinitionGenerateJsonSchema
+        DefinitionV11, schema_generator=ProjectDefinitionGenerateJsonSchema
     )["result"]
     for section in list_of_sections:
         _render_definition_description(root, section)

snowflake/cli/app/loggers.py

@@ -20,9 +20,6 @@ from dataclasses import asdict, dataclass, field
 from typing import Any, Dict, List
 
 import typer
-from snowflake.cli.api.config import (
-    get_logs_config,
-)
 from snowflake.cli.api.exceptions import InvalidLogsConfiguration
 from snowflake.cli.api.secure_path import SecurePath
 from snowflake.connector.errors import ConfigSourceError
@@ -103,6 +100,10 @@ def _remove_underscore_prefixes_from_keys(d: Dict[str, Any]) -> None:
 
 class FileLogsConfig:
     def __init__(self, debug: bool) -> None:
+        from snowflake.cli.api.config import (
+            get_logs_config,
+        )
+
         config = get_logs_config()
 
         self.path: SecurePath = SecurePath(config["path"])
@@ -142,8 +143,9 @@ def create_initial_loggers():
     config = InitialLoggingConfig()
     try:
         file_logs_config = FileLogsConfig(debug=False)
-        config.handlers["file"]["filename"] = file_logs_config.filename
-        _configurate_logging(config)
+        if file_logs_config.save_logs:
+            config.handlers["file"]["filename"] = file_logs_config.filename
+            _configurate_logging(config)
     except ConfigSourceError:
         pass
 
@@ -181,7 +183,9 @@ def create_loggers(verbose: bool, debug: bool):
     else:
         # We need to remove handler definition - otherwise it creates file even if `save_logs` is False
         del config.handlers["file"]
-        config.loggers["snowflake.cli"].handlers.remove("file")
+        for logger in config.loggers.values():
+            if "file" in logger.handlers:
+                logger.handlers.remove("file")
 
     config.loggers["snowflake.cli"].level = global_log_level
     config.loggers["snowflake"].level = global_log_level

snowflake/cli/app/printing.py

@@ -34,6 +34,7 @@ from snowflake.cli.api.output.types import (
     MessageResult,
     MultipleResults,
     ObjectResult,
+    StreamResult,
 )
 from snowflake.cli.api.sanitizers import sanitize_for_terminal
 
@@ -89,7 +90,7 @@ def _print_multiple_table_results(obj: CollectionResult):
         for item in items:
             table.add_row(*[str(i) for i in item.values()])
     # Add separator between tables
-    rich_print()
+    rich_print(flush=True)
 
 
 def is_structured_format(output_format):
@@ -98,12 +99,21 @@ def is_structured_format(output_format):
 
 def print_structured(result: CommandResult):
     """Handles outputs like json, yml and other structured and parsable formats."""
+    printed_end_line = False
     if isinstance(result, MultipleResults):
         _stream_json(result)
+    elif isinstance(result, StreamResult):
+        # A StreamResult prints each value onto its own line
+        # instead of joining all the values into a JSON array
+        for r in result.result:
+            json.dump(r, sys.stdout, cls=CustomJSONEncoder)
+            print(flush=True)
+            printed_end_line = True
     else:
         json.dump(result, sys.stdout, cls=CustomJSONEncoder, indent=4)
     # Adds empty line at the end
-    print()
+    if not printed_end_line:
+        print(flush=True)
 
 
 def _stream_json(result):
@@ -130,11 +140,11 @@ def _stream_json(result):
 def print_unstructured(obj: CommandResult | None):
     """Handles outputs like table, plain text and other unstructured types."""
     if not obj:
-        rich_print("Done")
+        rich_print("Done", flush=True)
     elif not obj.result:
-        rich_print("No data")
+        rich_print("No data", flush=True)
     elif isinstance(obj, MessageResult):
-        rich_print(sanitize_for_terminal(obj.message))
+        rich_print(sanitize_for_terminal(obj.message), flush=True)
     else:
         if isinstance(obj, ObjectResult):
             _print_single_table(obj)
@@ -152,14 +162,14 @@ def _print_single_table(obj):
         table.add_row(
             sanitize_for_terminal(str(key)), sanitize_for_terminal(str(value))
         )
-    rich_print(table)
+    rich_print(table, flush=True)
 
 
 def print_result(cmd_result: CommandResult, output_format: OutputFormat | None = None):
     output_format = output_format or _get_format_type()
     if is_structured_format(output_format):
         print_structured(cmd_result)
-    elif isinstance(cmd_result, MultipleResults):
+    elif isinstance(cmd_result, (MultipleResults, StreamResult)):
         for res in cmd_result.result:
             print_result(res)
     elif (

snowflake/cli/app/snow_connector.py

@@ -21,7 +21,12 @@ from typing import Dict, Optional
 
 import snowflake.connector
 from click.exceptions import ClickException
-from snowflake.cli.api.config import get_connection_dict, get_default_connection_dict
+from snowflake.cli.api.cli_global_context import cli_context
+from snowflake.cli.api.config import (
+    get_connection_dict,
+    get_default_connection_dict,
+    get_default_connection_name,
+)
 from snowflake.cli.api.constants import DEFAULT_SIZE_LIMIT_MB
 from snowflake.cli.api.exceptions import (
     InvalidConnectionConfiguration,
@@ -70,6 +75,9 @@ def connect_to_snowflake(
         connection_parameters = {}  # we will apply overrides in next step
     else:
         connection_parameters = get_default_connection_dict()
+        cli_context.connection_context.set_connection_name(
+            get_default_connection_name()
+        )
 
     # Apply overrides to connection details
     for key, value in overrides.items():

snowflake/cli/app/telemetry.py

@@ -22,6 +22,7 @@ from typing import Any, Dict, Union
 import click
 from snowflake.cli.__about__ import VERSION
 from snowflake.cli.api.cli_global_context import cli_context
+from snowflake.cli.api.commands.execution_metadata import ExecutionMetadata
 from snowflake.cli.api.config import get_feature_flags_section
 from snowflake.cli.api.output.formats import OutputFormat
 from snowflake.cli.api.utils.error_handling import ignore_exceptions
@@ -44,19 +45,25 @@ class CLITelemetryField(Enum):
     COMMAND = "command"
     COMMAND_GROUP = "command_group"
     COMMAND_FLAGS = "command_flags"
+    COMMAND_EXECUTION_ID = "command_execution_id"
+    COMMAND_RESULT_STATUS = "command_result_status"
     COMMAND_OUTPUT_TYPE = "command_output_type"
+    COMMAND_EXECUTION_TIME = "command_execution_time"
     # Configuration
     CONFIG_FEATURE_FLAGS = "config_feature_flags"
     # Information
     EVENT = "event"
     ERROR_MSG = "error_msg"
     ERROR_TYPE = "error_type"
+    IS_CLI_EXCEPTION = "is_cli_exception"
     # Project context
     PROJECT_DEFINITION_VERSION = "project_definition_version"
 
 
 class TelemetryEvent(Enum):
     CMD_EXECUTION = "executing_command"
+    CMD_EXECUTION_ERROR = "error_executing_command"
+    CMD_EXECUTION_RESULT = "result_executing_command"
 
 
 TelemetryDict = Dict[Union[CLITelemetryField, TelemetryField], Any]
@@ -141,8 +148,40 @@ _telemetry = CLITelemetryClient(ctx=cli_context)
 
 
 @ignore_exceptions()
-def log_command_usage():
-    _telemetry.send({TelemetryField.KEY_TYPE: TelemetryEvent.CMD_EXECUTION.value})
+def log_command_usage(execution: ExecutionMetadata):
+    _telemetry.send(
+        {
+            TelemetryField.KEY_TYPE: TelemetryEvent.CMD_EXECUTION.value,
+            CLITelemetryField.COMMAND_EXECUTION_ID: execution.execution_id,
+        }
+    )
+
+
+@ignore_exceptions()
+def log_command_result(execution: ExecutionMetadata):
+    _telemetry.send(
+        {
+            TelemetryField.KEY_TYPE: TelemetryEvent.CMD_EXECUTION_RESULT.value,
+            CLITelemetryField.COMMAND_EXECUTION_ID: execution.execution_id,
+            CLITelemetryField.COMMAND_RESULT_STATUS: execution.status.value,
+            CLITelemetryField.COMMAND_EXECUTION_TIME: execution.get_duration(),
+        }
+    )
+
+
+@ignore_exceptions()
+def log_command_execution_error(exception: Exception, execution: ExecutionMetadata):
+    exception_type: str = type(exception).__name__
+    is_cli_exception: bool = issubclass(exception.__class__, click.ClickException)
+    _telemetry.send(
+        {
+            TelemetryField.KEY_TYPE: TelemetryEvent.CMD_EXECUTION_ERROR.value,
+            CLITelemetryField.COMMAND_EXECUTION_ID: execution.execution_id,
+            CLITelemetryField.ERROR_TYPE: exception_type,
+            CLITelemetryField.IS_CLI_EXCEPTION: is_cli_exception,
+            CLITelemetryField.COMMAND_EXECUTION_TIME: execution.get_duration(),
+        }
+    )
 
 
 @ignore_exceptions()

snowflake/cli/plugins/connection/commands.py

@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 import logging
+import os.path
 
 import typer
 from click import ClickException, Context, Parameter  # type: ignore
@@ -300,9 +301,9 @@ def test(
     }
 
     if conn_ctx.enable_diag:
-        result[
-            "Diag Report Location"
-        ] = f"{conn_ctx.diag_log_path}/SnowflakeConnectionTestReport.txt"
+        result["Diag Report Location"] = os.path.join(
+            conn_ctx.diag_log_path, "SnowflakeConnectionTestReport.txt"
+        )
 
     return ObjectResult(result)