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

snowflake/cli/{app → _app}/telemetry.py

@@ -21,12 +21,15 @@ 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._app.constants import PARAM_APPLICATION_NAME
+from snowflake.cli.api.cli_global_context import (
+    _CliGlobalContextAccess,
+    get_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
-from snowflake.cli.app.constants import PARAM_APPLICATION_NAME
 from snowflake.connector.telemetry import (
     TelemetryData,
     TelemetryField,
@@ -34,6 +37,12 @@ from snowflake.connector.telemetry import (
 from snowflake.connector.time_util import get_time_millis
 
 
+@unique
+class CLIInstallationSource(Enum):
+    BINARY = "binary"
+    PYPI = "pypi"
+
+
 @unique
 class CLITelemetryField(Enum):
     # Basic information
@@ -41,6 +50,7 @@ class CLITelemetryField(Enum):
     VERSION_CLI = "version_cli"
     VERSION_PYTHON = "version_python"
     VERSION_OS = "version_os"
+    INSTALLATION_SOURCE = "installation_source"
     # Command execution context
     COMMAND = "command"
     COMMAND_GROUP = "command_group"
@@ -51,6 +61,8 @@ class CLITelemetryField(Enum):
     COMMAND_EXECUTION_TIME = "command_execution_time"
     # Configuration
     CONFIG_FEATURE_FLAGS = "config_feature_flags"
+    # Metrics
+    COUNTERS = "counters"
     # Information
     EVENT = "event"
     ERROR_MSG = "error_msg"
@@ -69,6 +81,16 @@ class TelemetryEvent(Enum):
 TelemetryDict = Dict[Union[CLITelemetryField, TelemetryField], Any]
 
 
+def _get_command_metrics() -> TelemetryDict:
+    cli_context = get_cli_context()
+
+    return {
+        CLITelemetryField.COUNTERS: {
+            **cli_context.metrics.counters,
+        }
+    }
+
+
 def _find_command_info() -> TelemetryDict:
     ctx = click.get_current_context()
     command_path = ctx.command_path.split(" ")[1:]
@@ -88,13 +110,18 @@ def _find_command_info() -> TelemetryDict:
 
 
 def _get_definition_version() -> str | None:
-    from snowflake.cli.api.cli_global_context import cli_context
-
+    cli_context = get_cli_context()
     if cli_context.project_definition:
         return cli_context.project_definition.definition_version
     return None
 
 
+def _get_installation_source() -> CLIInstallationSource:
+    if getattr(sys, "frozen", False) and hasattr(sys, "_MEIPASS"):
+        return CLIInstallationSource.BINARY
+    return CLIInstallationSource.PYPI
+
+
 def command_info() -> str:
     info = _find_command_info()
     command = ".".join(info[CLITelemetryField.COMMAND])
@@ -107,8 +134,9 @@ def python_version() -> str:
 
 
 class CLITelemetryClient:
-    def __init__(self, ctx):
-        self._ctx = ctx
+    @property
+    def _ctx(self) -> _CliGlobalContextAccess:
+        return get_cli_context()
 
     @staticmethod
     def generate_telemetry_data_dict(
@@ -116,6 +144,7 @@ class CLITelemetryClient:
     ) -> Dict[str, Any]:
         data = {
             CLITelemetryField.SOURCE: PARAM_APPLICATION_NAME,
+            CLITelemetryField.INSTALLATION_SOURCE: _get_installation_source().value,
             CLITelemetryField.VERSION_CLI: VERSION,
             CLITelemetryField.VERSION_OS: platform.platform(),
             CLITelemetryField.VERSION_PYTHON: python_version(),
@@ -144,7 +173,7 @@ class CLITelemetryClient:
         self._telemetry.send_batch()
 
 
-_telemetry = CLITelemetryClient(ctx=cli_context)
+_telemetry = CLITelemetryClient()
 
 
 @ignore_exceptions()
@@ -165,6 +194,7 @@ def log_command_result(execution: ExecutionMetadata):
             CLITelemetryField.COMMAND_EXECUTION_ID: execution.execution_id,
             CLITelemetryField.COMMAND_RESULT_STATUS: execution.status.value,
             CLITelemetryField.COMMAND_EXECUTION_TIME: execution.get_duration(),
+            **_get_command_metrics(),
         }
     )
 
@@ -180,6 +210,7 @@ def log_command_execution_error(exception: Exception, execution: ExecutionMetada
             CLITelemetryField.ERROR_TYPE: exception_type,
             CLITelemetryField.IS_CLI_EXCEPTION: is_cli_exception,
             CLITelemetryField.COMMAND_EXECUTION_TIME: execution.get_duration(),
+            **_get_command_metrics(),
         }
     )
 

snowflake/cli/_app/version_check.py

@@ -0,0 +1,74 @@
+import json
+import time
+
+import requests
+from packaging.version import Version
+from snowflake.cli.__about__ import VERSION
+from snowflake.cli.api.console import cli_console
+from snowflake.cli.api.secure_path import SecurePath
+from snowflake.connector.config_manager import CONFIG_MANAGER
+
+
+def get_new_version_msg() -> str | None:
+    last = _VersionCache().get_last_version()
+    current = Version(VERSION)
+    if last and last > current:
+        return f"\nNew version of Snowflake CLI available. Newest: {last}, current: {VERSION}\n"
+    return None
+
+
+def show_new_version_banner_callback(msg):
+    def _callback(*args, **kwargs):
+        if msg:
+            cli_console.message(msg)
+
+    return _callback
+
+
+class _VersionCache:
+    _last_time = "last_time_check"
+    _version = "version"
+    _version_cache_file = SecurePath(
+        CONFIG_MANAGER.file_path.parent / ".cli_version.cache"
+    )
+
+    def __init__(self):
+        self._cache_file = _VersionCache._version_cache_file
+
+    def _save_latest_version(self, version: str):
+        data = {
+            _VersionCache._last_time: time.time(),
+            _VersionCache._version: str(version),
+        }
+        self._cache_file.write_text(json.dumps(data))
+
+    @staticmethod
+    def _get_version_from_pypi() -> str | None:
+        headers = {"Content-Type": "application/vnd.pypi.simple.v1+json"}
+        response = requests.get(
+            "https://pypi.org/pypi/snowflake-cli-labs/json", headers=headers, timeout=3
+        )
+        response.raise_for_status()
+        return response.json()["info"]["version"]
+
+    def _update_latest_version(self) -> Version | None:
+        version = self._get_version_from_pypi()
+        if version is None:
+            return None
+        self._save_latest_version(version)
+        return Version(version)
+
+    def _read_latest_version(self) -> Version | None:
+        if self._cache_file.exists():
+            data = json.loads(self._cache_file.read_text())
+            now = time.time()
+            if data[_VersionCache._last_time] > now - 60 * 60:
+                return Version(data[_VersionCache._version])
+
+        return self._update_latest_version()
+
+    def get_last_version(self) -> Version | None:
+        try:
+            return self._read_latest_version()
+        except:  # anything, this it not crucial feature
+            return None

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

@@ -21,14 +21,19 @@ import typer
 from click import ClickException, Context, Parameter  # type: ignore
 from click.core import ParameterSource  # type: ignore
 from click.types import StringParamType
-from snowflake.cli.api.cli_global_context import cli_context
+from snowflake.cli._plugins.connection.util import (
+    strip_and_check_if_exists,
+    strip_if_value_present,
+)
+from snowflake.cli._plugins.object.manager import ObjectManager
+from snowflake.cli.api.cli_global_context import get_cli_context
 from snowflake.cli.api.commands.flags import (
     PLAIN_PASSWORD_MSG,
 )
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
 from snowflake.cli.api.config import (
     ConnectionConfig,
-    add_connection,
+    add_connection_to_proper_file,
     connection_exists,
     get_all_connections,
     get_connection_dict,
@@ -43,9 +48,8 @@ from snowflake.cli.api.output.types import (
     MessageResult,
     ObjectResult,
 )
-from snowflake.cli.plugins.object.manager import ObjectManager
 from snowflake.connector import ProgrammingError
-from snowflake.connector.config_manager import CONFIG_MANAGER
+from snowflake.connector.constants import CONNECTIONS_FILE
 
 app = SnowTyperFactory(
     name="connection",
@@ -87,6 +91,11 @@ def list_connections(**options) -> CommandResult:
         }
         for connection_name, connection_config in connections.items()
     )
+
+    if CONNECTIONS_FILE.exists():
+        cli_console.warning(
+            f"Reading connections from {CONNECTIONS_FILE}. Entries from config.toml are ignored."
+        )
     return CollectionResult(result)
 
 
@@ -94,8 +103,8 @@ def require_integer(field_name: str):
     def callback(value: str):
         if value is None:
             return None
-        if value.isdigit():
-            return value
+        if value.strip().isdigit():
+            return value.strip()
         raise ClickException(f"Value of {field_name} must be integer")
 
     return callback
@@ -117,6 +126,7 @@ def add(
         prompt="Name for this connection",
         help="Name of the new connection.",
         show_default=False,
+        callback=strip_if_value_present,
     ),
     account: str = typer.Option(
         None,
@@ -126,6 +136,7 @@ def add(
         prompt="Snowflake account name",
         help="Account name to use when authenticating with Snowflake.",
         show_default=False,
+        callback=strip_if_value_present,
     ),
     user: str = typer.Option(
         None,
@@ -135,6 +146,7 @@ def add(
         prompt="Snowflake username",
         show_default=False,
         help="Username to connect to Snowflake.",
+        callback=strip_if_value_present,
     ),
     password: str = typer.Option(
         EmptyInput(),
@@ -153,6 +165,7 @@ def add(
         click_type=OptionalPrompt(),
         prompt="Role for the connection",
         help="Role to use on Snowflake.",
+        callback=strip_if_value_present,
     ),
     warehouse: str = typer.Option(
         EmptyInput(),
@@ -161,6 +174,7 @@ def add(
         click_type=OptionalPrompt(),
         prompt="Warehouse for the connection",
         help="Warehouse to use on Snowflake.",
+        callback=strip_if_value_present,
     ),
     database: str = typer.Option(
         EmptyInput(),
@@ -169,6 +183,7 @@ def add(
         click_type=OptionalPrompt(),
         prompt="Database for the connection",
         help="Database to use on Snowflake.",
+        callback=strip_if_value_present,
     ),
     schema: str = typer.Option(
         EmptyInput(),
@@ -177,6 +192,7 @@ def add(
         click_type=OptionalPrompt(),
         prompt="Schema for the connection",
         help="Schema to use on Snowflake.",
+        callback=strip_if_value_present,
     ),
     host: str = typer.Option(
         EmptyInput(),
@@ -185,6 +201,7 @@ def add(
         click_type=OptionalPrompt(),
         prompt="Connection host",
         help="Host name the connection attempts to connect to Snowflake.",
+        callback=strip_if_value_present,
     ),
     port: int = typer.Option(
         EmptyInput(),
@@ -202,6 +219,7 @@ def add(
         click_type=OptionalPrompt(),
         prompt="Snowflake region",
         help="Region name if not the default Snowflake deployment.",
+        callback=strip_if_value_present,
     ),
     authenticator: str = typer.Option(
         EmptyInput(),
@@ -211,13 +229,15 @@ def add(
         prompt="Authentication method",
         help="Chosen authenticator, if other than password-based",
     ),
-    private_key_path: str = typer.Option(
+    private_key_file: str = typer.Option(
         EmptyInput(),
         "--private-key",
+        "--private-key-path",
         "-k",
         click_type=OptionalPrompt(),
         prompt="Path to private key file",
         help="Path to file containing private key",
+        callback=strip_and_check_if_exists,
     ),
     token_file_path: str = typer.Option(
         EmptyInput(),
@@ -226,6 +246,7 @@ def add(
         click_type=OptionalPrompt(),
         prompt="Path to token file",
         help="Path to file with an OAuth token that should be used when connecting to Snowflake",
+        callback=strip_and_check_if_exists,
     ),
     set_as_default: bool = typer.Option(
         False,
@@ -239,7 +260,7 @@ def add(
     if connection_exists(connection_name):
         raise ClickException(f"Connection {connection_name} already exists")
 
-    add_connection(
+    connections_file = add_connection_to_proper_file(
         connection_name,
         ConnectionConfig(
             account=account,
@@ -253,7 +274,7 @@ def add(
             warehouse=warehouse,
             role=role,
             authenticator=authenticator,
-            private_key_path=private_key_path,
+            private_key_file=private_key_file,
             token_file_path=token_file_path,
         ),
     )
@@ -263,7 +284,7 @@ def add(
         )
 
     return MessageResult(
-        f"Wrote new connection {connection_name} to {CONFIG_MANAGER.file_path}"
+        f"Wrote new connection {connection_name} to {connections_file}"
     )
 
 
@@ -276,6 +297,7 @@ def test(
     """
 
     # Test connection
+    cli_context = get_cli_context()
     conn = cli_context.connection
 
     # Test session attributes
@@ -320,7 +342,8 @@ def test(
 @app.command(requires_connection=False)
 def set_default(
     name: str = typer.Argument(
-        help="Name of the connection, as defined in your `config.toml`"
+        help="Name of the connection, as defined in your `config.toml`",
+        show_default=False,
     ),
     **options,
 ):

snowflake/cli/_plugins/connection/plugin_spec.py

@@ -0,0 +1,30 @@
+# 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 snowflake.cli._plugins.connection import commands
+from snowflake.cli.api.plugins.command import (
+    SNOWCLI_ROOT_COMMAND_PATH,
+    CommandSpec,
+    CommandType,
+    plugin_hook_impl,
+)
+
+
+@plugin_hook_impl
+def command_spec():
+    return CommandSpec(
+        parent_command_path=SNOWCLI_ROOT_COMMAND_PATH,
+        command_type=CommandType.COMMAND_GROUP,
+        typer_instance=commands.app.create_instance(),
+    )

snowflake/cli/{plugins → _plugins}/connection/util.py

@@ -16,6 +16,8 @@ from __future__ import annotations
 
 import json
 import logging
+import os
+from typing import Optional
 
 from click.exceptions import ClickException
 from snowflake.connector import SnowflakeConnection
@@ -177,3 +179,17 @@ def make_snowsight_url(conn: SnowflakeConnection, path: str) -> str:
     account = get_account(conn)
     path_with_slash = path if path.startswith("/") else f"/{path}"
     return f"{snowsight_host}/{deployment}/{account}{path_with_slash}"
+
+
+def strip_if_value_present(value: Optional[str]) -> Optional[str]:
+    return value.strip() if value else value
+
+
+def ensure_that_path_exist(path: Optional[str]) -> Optional[str]:
+    if path and not os.path.exists(path):
+        raise ClickException(f"Path {path} does not exist.")
+    return path
+
+
+def strip_and_check_if_exists(value: Optional[str]) -> Optional[str]:
+    return ensure_that_path_exist(strip_if_value_present(value))

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

@@ -21,8 +21,20 @@ from typing import List, Optional
 import click
 import typer
 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._plugins.cortex.constants import DEFAULT_MODEL
+from snowflake.cli._plugins.cortex.manager import CortexManager
+from snowflake.cli._plugins.cortex.types import (
+    Language,
+    Model,
+    Question,
+    SourceDocument,
+    Text,
+)
+from snowflake.cli.api.cli_global_context import get_cli_context
+from snowflake.cli.api.commands.overrideable_parameter import (
+    OverrideableArgument,
+    OverrideableOption,
+)
 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 (
@@ -31,15 +43,6 @@ from snowflake.cli.api.output.types import (
     MessageResult,
 )
 from snowflake.cli.api.secure_path import SecurePath
-from snowflake.cli.plugins.cortex.constants import DEFAULT_MODEL
-from snowflake.cli.plugins.cortex.manager import CortexManager
-from snowflake.cli.plugins.cortex.types import (
-    Language,
-    Model,
-    Question,
-    SourceDocument,
-    Text,
-)
 
 app = SnowTyperFactory(
     name="cortex",
@@ -48,6 +51,25 @@ app = SnowTyperFactory(
 
 SEARCH_COMMAND_ENABLED = sys.version_info < PYTHON_3_12
 
+SOURCE_EXCLUSIVE_OPTION_NAMES = ["text", "file", "source_document_text"]
+
+# Creates a Typer option and verifies if the mutually exclusive options are set in the command.
+ExclusiveReadableFileOption = OverrideableOption(
+    None,
+    "--file",
+    mutually_exclusive=SOURCE_EXCLUSIVE_OPTION_NAMES,
+    exists=True,
+    file_okay=True,
+    dir_okay=False,
+    readable=True,
+    show_default=False,
+)
+
+# Creates a Typer argument and verifies if the mutually exclusive options are set in the command.
+ExclusiveTextSourceArgument = OverrideableArgument(
+    mutually_exclusive=SOURCE_EXCLUSIVE_OPTION_NAMES,
+)
+
 
 @app.command(
     requires_connection=True,
@@ -79,7 +101,7 @@ def search(
     if not columns:
         columns = []
 
-    conn = cli_context.connection
+    conn = get_cli_context().connection
 
     search_service = (
         Root(conn)
@@ -100,8 +122,8 @@ def search(
     requires_connection=True,
 )
 def complete(
-    text: Optional[str] = typer.Argument(
-        None,
+    text: Optional[str] = ExclusiveTextSourceArgument(
+        default=None,
         help="Prompt to be used to generate a completion. Cannot be combined with --file option.",
         show_default=False,
     ),
@@ -110,9 +132,8 @@ def complete(
         "--model",
         help="String specifying the model to be used.",
     ),
-    file: Optional[Path] = readable_file_option(
-        param_name="--file",
-        help_str="JSON file containing conversation history to be used to generate a completion. Cannot be combined with TEXT argument.",
+    file: Optional[Path] = ExclusiveReadableFileOption(
+        help="JSON file containing conversation history to be used to generate a completion. Cannot be combined with TEXT argument.",
     ),
     **options,
 ) -> CommandResult:
@@ -124,8 +145,6 @@ def complete(
 
     manager = CortexManager()
 
-    if text and file:
-        raise UsageError("--file option cannot be used together with TEXT argument.")
     if text:
         result_text = manager.complete_for_prompt(
             text=Text(text),
@@ -152,14 +171,13 @@ def extract_answer(
         help="String containing the question to be answered.",
         show_default=False,
     ),
-    source_document_text: Optional[str] = typer.Argument(
-        None,
+    source_document_text: Optional[str] = ExclusiveTextSourceArgument(
+        default=None,
         help="String containing the plain-text or JSON document that contains the answer to the question. Cannot be combined with --file option.",
         show_default=False,
     ),
-    file: Optional[Path] = readable_file_option(
-        param_name="--file",
-        help_str="File containing the plain-text or JSON document that contains the answer to the question. Cannot be combined with SOURCE_DOCUMENT_TEXT argument.",
+    file: Optional[Path] = ExclusiveReadableFileOption(
+        help="File containing the plain-text or JSON document that contains the answer to the question. Cannot be combined with SOURCE_DOCUMENT_TEXT argument.",
     ),
     **options,
 ) -> CommandResult:
@@ -170,10 +188,6 @@ def extract_answer(
 
     manager = CortexManager()
 
-    if source_document_text and file:
-        raise UsageError(
-            "--file option cannot be used together with SOURCE_DOCUMENT_TEXT argument."
-        )
     if source_document_text:
         result_text = manager.extract_answer_from_source_document(
             source_document=SourceDocument(source_document_text),
@@ -197,14 +211,13 @@ def extract_answer(
     requires_connection=True,
 )
 def sentiment(
-    text: Optional[str] = typer.Argument(
-        None,
+    text: Optional[str] = ExclusiveTextSourceArgument(
+        default=None,
         help="String containing the text for which a sentiment score should be calculated. Cannot be combined with --file option.",
         show_default=False,
     ),
-    file: Optional[Path] = readable_file_option(
-        param_name="--file",
-        help_str="File containing the text for which a sentiment score should be calculated. Cannot be combined with TEXT argument.",
+    file: Optional[Path] = ExclusiveReadableFileOption(
+        help="File containing the text for which a sentiment score should be calculated. Cannot be combined with TEXT argument.",
     ),
     **options,
 ) -> CommandResult:
@@ -216,8 +229,6 @@ def sentiment(
 
     manager = CortexManager()
 
-    if text and file:
-        raise UsageError("--file option cannot be used together with TEXT argument.")
     if text:
         result_text = manager.calculate_sentiment_for_text(
             text=Text(text),
@@ -237,14 +248,13 @@ def sentiment(
     requires_connection=True,
 )
 def summarize(
-    text: Optional[str] = typer.Argument(
-        None,
+    text: Optional[str] = ExclusiveTextSourceArgument(
+        default=None,
         help="String containing the English text from which a summary should be generated. Cannot be combined with --file option.",
         show_default=False,
     ),
-    file: Optional[Path] = readable_file_option(
-        param_name="--file",
-        help_str="File containing the English text from which a summary should be generated. Cannot be combined with TEXT argument.",
+    file: Optional[Path] = ExclusiveReadableFileOption(
+        help="File containing the English text from which a summary should be generated. Cannot be combined with TEXT argument.",
     ),
     **options,
 ) -> CommandResult:
@@ -254,8 +264,6 @@ def summarize(
 
     manager = CortexManager()
 
-    if text and file:
-        raise UsageError("--file option cannot be used together with TEXT argument.")
     if text:
         result_text = manager.summarize_text(
             text=Text(text),
@@ -275,8 +283,8 @@ def summarize(
     requires_connection=True,
 )
 def translate(
-    text: Optional[str] = typer.Argument(
-        None,
+    text: Optional[str] = ExclusiveTextSourceArgument(
+        default=None,
         help="String containing the text to be translated. Cannot be combined with --file option.",
         show_default=False,
     ),
@@ -292,9 +300,8 @@ def translate(
         help="String specifying the language code into which the text should be translated. See Snowflake Cortex documentation for a list of supported language codes.",
         show_default=False,
     ),
-    file: Optional[Path] = readable_file_option(
-        param_name="--file",
-        help_str="File containing the text to be translated. Cannot be combined with TEXT argument.",
+    file: Optional[Path] = ExclusiveReadableFileOption(
+        help="File containing the text to be translated. Cannot be combined with TEXT argument.",
     ),
     **options,
 ) -> CommandResult:
@@ -307,8 +314,6 @@ def translate(
     source_language = None if from_language is None else Language(from_language)
     target_language = Language(to_language)
 
-    if text and file:
-        raise UsageError("--file option cannot be used together with TEXT argument.")
     if text:
         result_text = manager.translate_text(
             text=Text(text),

snowflake/cli/{plugins → _plugins}/cortex/constants.py

@@ -12,6 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from snowflake.cli.plugins.cortex.types import Model
+from snowflake.cli._plugins.cortex.types import Model
 
 DEFAULT_MODEL: Model = Model("snowflake-arctic")