snowflake-cli-labs 2.6.1__py3-none-any.whl → 2.7.0rc1__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/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/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

@@ -183,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)
 

snowflake/cli/plugins/connection/util.py

@@ -12,14 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
+import json
 import logging
 
 from click.exceptions import ClickException
 from snowflake.connector import SnowflakeConnection
 from snowflake.connector.cursor import DictCursor
 
-LOCAL_DEPLOYMENT: str = "us-west-2"
-
 log = logging.getLogger(__name__)
 
 REGIONLESS_QUERY = """
@@ -29,11 +30,23 @@ REGIONLESS_QUERY = """
     )) where value['name'] = 'UI_SNOWSIGHT_ENABLE_REGIONLESS_REDIRECT';
 """
 
+ALLOWLIST_QUERY = "SELECT SYSTEM$ALLOWLIST()"
+SNOWFLAKE_DEPLOYMENT = "SNOWFLAKE_DEPLOYMENT"
+LOCAL_DEPLOYMENT_REGION: str = "us-west-2"
 
-class MissingConnectionHostError(ClickException):
+
+class MissingConnectionAccountError(ClickException):
     def __init__(self, conn: SnowflakeConnection):
         super().__init__(
-            f"The connection host ({conn.host}) was missing or not in "
+            "Could not determine account by system call, configured account name, or configured host. Connection: "
+            + repr(conn)
+        )
+
+
+class MissingConnectionRegionError(ClickException):
+    def __init__(self, host: str | None):
+        super().__init__(
+            f"The connection host ({host}) was missing or not in "
             "the expected format "
             "(<account>.<deployment>.snowflakecomputing.com)"
         )
@@ -50,11 +63,60 @@ def is_regionless_redirect(conn: SnowflakeConnection) -> bool:
         *_, cursor = conn.execute_string(REGIONLESS_QUERY, cursor_class=DictCursor)
         return cursor.fetchone()["REGIONLESS"].lower() == "true"
     except:
-        # by default, assume that
-        log.warning("Cannot determine regionless redirect; assuming True.")
+        log.warning(
+            "Cannot determine regionless redirect; assuming True.", exc_info=True
+        )
         return True
 
 
+def get_host_region(host: str) -> str | None:
+    """
+    Looks for hosts of form
+    <account>.[x.y.z].snowflakecomputing.com
+    Returns the three-part [region identifier] or None.
+    """
+    host_parts = host.split(".")
+    if host_parts[-1] == "local":
+        return LOCAL_DEPLOYMENT_REGION
+    elif len(host_parts) == 6:
+        return ".".join(host_parts[1:4])
+    return None
+
+
+def guess_regioned_host_from_allowlist(conn: SnowflakeConnection) -> str | None:
+    """
+    Use SYSTEM$ALLOWLIST to find a regioned host (<account>.x.y.z.snowflakecomputing.com)
+    that corresponds to the given Snowflake connection object.
+    """
+    try:
+        *_, cursor = conn.execute_string(ALLOWLIST_QUERY, cursor_class=DictCursor)
+        allowlist_tuples = json.loads(cursor.fetchone()["SYSTEM$ALLOWLIST()"])
+        for t in allowlist_tuples:
+            if t["type"] == SNOWFLAKE_DEPLOYMENT:
+                if get_host_region(t["host"]) is not None:
+                    return t["host"]
+    except:
+        log.warning(
+            "Could not call SYSTEM$ALLOWLIST; returning an empty guess.", exc_info=True
+        )
+    return None
+
+
+def get_region(conn: SnowflakeConnection) -> str:
+    """
+    Get the region of the given connection, or raise MissingConnectionRegionError.
+    """
+    if conn.host:
+        if region := get_host_region(conn.host):
+            return region
+
+    if host := guess_regioned_host_from_allowlist(conn):
+        if region := get_host_region(host):
+            return region
+
+    raise MissingConnectionRegionError(host or conn.host)
+
+
 def get_context(conn: SnowflakeConnection) -> str:
     """
     Determines the first part of the path in a Snowsight URL.
@@ -67,14 +129,7 @@ def get_context(conn: SnowflakeConnection) -> str:
         )
         return cursor.fetchone()["SYSTEM$RETURN_CURRENT_ORG_NAME()"]
 
-    host_parts = conn.host.split(".")
-    if host_parts[-1] == "local":
-        return LOCAL_DEPLOYMENT
-
-    if len(host_parts) == 6:
-        return ".".join(host_parts[1:4])
-
-    raise MissingConnectionHostError(conn)
+    return get_region(conn)
 
 
 def get_account(conn: SnowflakeConnection) -> str:
@@ -91,11 +146,11 @@ def get_account(conn: SnowflakeConnection) -> str:
         if conn.account:
             return conn.account
 
-        if not conn.host:
-            raise MissingConnectionHostError(conn)
+        if conn.host:
+            host_parts = conn.host.split(".")
+            return host_parts[0]
 
-        host_parts = conn.host.split(".")
-        return host_parts[0]
+        raise MissingConnectionAccountError(conn)
 
 
 def get_snowsight_host(conn: SnowflakeConnection) -> str:

snowflake/cli/plugins/cortex/commands.py

@@ -24,6 +24,7 @@ from click import UsageError
 from snowflake.cli.api.cli_global_context import cli_context
 from snowflake.cli.api.commands.flags import readable_file_option
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
+from snowflake.cli.api.constants import PYTHON_3_12
 from snowflake.cli.api.output.types import (
     CollectionResult,
     CommandResult,
@@ -45,7 +46,7 @@ app = SnowTyperFactory(
     help="Provides access to Snowflake Cortex.",
 )
 
-SEARCH_COMMAND_ENABLED = sys.version_info < (3, 12)
+SEARCH_COMMAND_ENABLED = sys.version_info < PYTHON_3_12
 
 
 @app.command(

snowflake/cli/plugins/git/commands.py

@@ -14,15 +14,18 @@
 
 from __future__ import annotations
 
+import itertools
 import logging
+from os import path
+from pathlib import Path
 from typing import List, Optional
 
 import typer
 from click import ClickException
 from snowflake.cli.api.commands.flags import (
+    ExecuteVariablesOption,
     OnErrorOption,
     PatternOption,
-    VariablesOption,
     identifier_argument,
     like_option,
 )
@@ -37,7 +40,6 @@ from snowflake.cli.plugins.object.command_aliases import (
     scope_option,
 )
 from snowflake.cli.plugins.object.manager import ObjectManager
-from snowflake.cli.plugins.stage.commands import get
 from snowflake.cli.plugins.stage.manager import OnErrorType
 
 app = SnowTyperFactory(
@@ -264,7 +266,6 @@ def copy(
             )
         )
     return get(
-        recursive=True,
         source_path=repository_path,
         destination_path=destination_path,
         parallel=parallel,
@@ -275,7 +276,7 @@ def copy(
 def execute(
     repository_path: str = RepoPathArgument,
     on_error: OnErrorType = OnErrorOption,
-    variables: Optional[List[str]] = VariablesOption,
+    variables: Optional[List[str]] = ExecuteVariablesOption,
     **options,
 ):
     """
@@ -287,3 +288,18 @@ def execute(
         stage_path=repository_path, on_error=on_error, variables=variables
     )
     return CollectionResult(results)
+
+
+def get(source_path: str, destination_path: str, parallel: int):
+    target = Path(destination_path).resolve()
+
+    cursors = GitManager().get_recursive(
+        stage_path=source_path, dest_path=target, parallel=parallel
+    )
+    results = [list(QueryResult(c).result) for c in cursors]
+    flattened_results = list(itertools.chain.from_iterable(results))
+    sorted_results = sorted(
+        flattened_results,
+        key=lambda e: (path.dirname(e["file"]), path.basename(e["file"])),
+    )
+    return CollectionResult(sorted_results)