snowflake-cli 3.2.2__py3-none-any.whl → 3.4.1__py3-none-any.whl

snowflake/cli/_plugins/cortex/commands.py

@@ -96,16 +96,14 @@ def search(
             "Cortex Search uses Snowflake Python API that currently does not support your Python version"
         )
 
-    from snowflake.core import Root
-
     if not columns:
         columns = []
 
     conn = get_cli_context().connection
+    root = get_cli_context().snow_api_root
 
     search_service = (
-        Root(conn)
-        .databases[conn.database]
+        root.databases[conn.database]
         .schemas[conn.schema]
         .cortex_search_services[service]
     )

snowflake/cli/_plugins/git/manager.py

@@ -116,7 +116,7 @@ class GitManager(StageManager):
         return f"{'/'.join(path_parts[ONLY_STAGE])}/"
 
     @staticmethod
-    def _stage_path_part_factory(stage_path: str) -> StagePathParts:
+    def stage_path_parts_from_str(stage_path: str) -> StagePathParts:
         stage_path = StageManager.get_standard_stage_prefix(stage_path)
         if stage_path.startswith(USER_STAGE_PREFIX):
             return UserStagePathParts(stage_path)

snowflake/cli/_plugins/helpers/commands.py

@@ -14,17 +14,30 @@
 
 from __future__ import annotations
 
+import logging
+from pathlib import Path
+from typing import Any, List, Optional
+
 import typer
 import yaml
 from click import ClickException
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
-from snowflake.cli.api.output.types import MessageResult
+from snowflake.cli.api.config import (
+    ConnectionConfig,
+    add_connection_to_proper_file,
+    get_all_connections,
+    set_config_value,
+)
+from snowflake.cli.api.console import cli_console
+from snowflake.cli.api.output.types import CommandResult, MessageResult
 from snowflake.cli.api.project.definition_conversion import (
     convert_project_definition_to_v2,
 )
 from snowflake.cli.api.project.definition_manager import DefinitionManager
 from snowflake.cli.api.secure_path import SecurePath
 
+log = logging.getLogger(__name__)
+
 app = SnowTyperFactory(
     name="helpers",
     help="Helper commands.",
@@ -88,3 +101,196 @@ def v1_to_v2(
             width=float("inf"),  # Don't break lines
         )
     return MessageResult("Project definition migrated to version 2.")
+
+
+@app.command(name="import-snowsql-connections", requires_connection=False)
+def import_snowsql_connections(
+    custom_snowsql_config_files: Optional[List[Path]] = typer.Option(
+        None,
+        "--snowsql-config-file",
+        help="Specifies file paths to custom SnowSQL configuration. The option can be used multiple times to specify more than 1 file.",
+        dir_okay=False,
+        exists=True,
+    ),
+    default_cli_connection_name: str = typer.Option(
+        "default",
+        "--default-connection-name",
+        help="Specifies the name which will be given in Snowflake CLI to the default connection imported from SnowSQL.",
+    ),
+    **options,
+) -> CommandResult:
+    """Import your existing connections from your SnowSQL configuration."""
+
+    snowsql_config_files: list[Path] = custom_snowsql_config_files or [
+        Path("/etc/snowsql.cnf"),
+        Path("/etc/snowflake/snowsql.cnf"),
+        Path("/usr/local/etc/snowsql.cnf"),
+        Path.home() / Path(".snowsql.cnf"),
+        Path.home() / Path(".snowsql/config"),
+    ]
+    snowsql_config_secure_paths: list[SecurePath] = [
+        SecurePath(p) for p in snowsql_config_files
+    ]
+
+    all_imported_connections = _read_all_connections_from_snowsql(
+        default_cli_connection_name, snowsql_config_secure_paths
+    )
+    _validate_and_save_connections_imported_from_snowsql(
+        default_cli_connection_name, all_imported_connections
+    )
+    return MessageResult(
+        "Connections successfully imported from SnowSQL to Snowflake CLI."
+    )
+
+
+def _read_all_connections_from_snowsql(
+    default_cli_connection_name: str, snowsql_config_files: List[SecurePath]
+) -> dict[str, dict]:
+    import configparser
+
+    imported_default_connection: dict[str, Any] = {}
+    imported_named_connections: dict[str, dict] = {}
+
+    for file in snowsql_config_files:
+        if not file.exists():
+            cli_console.step(
+                f"SnowSQL config file [{str(file.path)}] does not exist. Skipping."
+            )
+            continue
+
+        cli_console.step(f"Trying to read connections from [{str(file.path)}].")
+        snowsql_config = configparser.ConfigParser()
+        snowsql_config.read(file.path)
+
+        if "connections" in snowsql_config and snowsql_config.items("connections"):
+            cli_console.step(
+                f"Reading SnowSQL's default connection configuration from [{str(file.path)}]"
+            )
+            snowsql_default_connection = snowsql_config.items("connections")
+            imported_default_connection.update(
+                _convert_connection_from_snowsql_config_section(
+                    snowsql_default_connection
+                )
+            )
+
+        other_snowsql_connection_section_names = [
+            section_name
+            for section_name in snowsql_config.sections()
+            if section_name.startswith("connections.")
+        ]
+        for snowsql_connection_section_name in other_snowsql_connection_section_names:
+            cli_console.step(
+                f"Reading SnowSQL's connection configuration [{snowsql_connection_section_name}] from [{str(file.path)}]"
+            )
+            snowsql_named_connection = snowsql_config.items(
+                snowsql_connection_section_name
+            )
+            if not snowsql_named_connection:
+                cli_console.step(
+                    f"Empty connection configuration [{snowsql_connection_section_name}] in [{str(file.path)}]. Skipping."
+                )
+                continue
+
+            connection_name = snowsql_connection_section_name.removeprefix(
+                "connections."
+            )
+            imported_named_conenction = _convert_connection_from_snowsql_config_section(
+                snowsql_named_connection
+            )
+            if connection_name in imported_named_connections:
+                imported_named_connections[connection_name].update(
+                    imported_named_conenction
+                )
+            else:
+                imported_named_connections[connection_name] = imported_named_conenction
+
+    def imported_default_connection_as_named_connection():
+        name = _validate_imported_default_connection_name(
+            default_cli_connection_name, imported_named_connections
+        )
+        return {name: imported_default_connection}
+
+    named_default_connection = (
+        imported_default_connection_as_named_connection()
+        if imported_default_connection
+        else {}
+    )
+
+    return imported_named_connections | named_default_connection
+
+
+def _validate_imported_default_connection_name(
+    name_candidate: str, other_snowsql_connections: dict[str, dict]
+) -> str:
+    if name_candidate in other_snowsql_connections:
+        new_name_candidate = typer.prompt(
+            f"Chosen default connection name '{name_candidate}' is already taken by other connection being imported from SnowSQL. Please choose a different name for your default connection"
+        )
+        return _validate_imported_default_connection_name(
+            new_name_candidate, other_snowsql_connections
+        )
+    else:
+        return name_candidate
+
+
+def _convert_connection_from_snowsql_config_section(
+    snowsql_connection: list[tuple[str, Any]]
+) -> dict[str, Any]:
+    from ast import literal_eval
+
+    key_names_replacements = {
+        "accountname": "account",
+        "username": "user",
+        "databasename": "database",
+        "dbname": "database",
+        "schemaname": "schema",
+        "warehousename": "warehouse",
+        "rolename": "role",
+        "private_key_path": "private_key_file",
+    }
+
+    def parse_value(value: Any):
+        try:
+            parsed_value = literal_eval(value)
+        except Exception:
+            parsed_value = value
+        return parsed_value
+
+    cli_connection: dict[str, Any] = {}
+    for key, value in snowsql_connection:
+        cli_key = key_names_replacements.get(key, key)
+        cli_value = parse_value(value)
+        cli_connection[cli_key] = cli_value
+    return cli_connection
+
+
+def _validate_and_save_connections_imported_from_snowsql(
+    default_cli_connection_name: str, all_imported_connections: dict[str, Any]
+):
+    existing_cli_connection_names: set[str] = set(get_all_connections().keys())
+    imported_connections_to_save: dict[str, Any] = {}
+    for (
+        imported_connection_name,
+        imported_connection,
+    ) in all_imported_connections.items():
+        if imported_connection_name in existing_cli_connection_names:
+            override_cli_connection = typer.confirm(
+                f"Connection '{imported_connection_name}' already exists in Snowflake CLI, do you want to use SnowSQL definition and override existing connection in Snowflake CLI?"
+            )
+            if not override_cli_connection:
+                continue
+        imported_connections_to_save[imported_connection_name] = imported_connection
+
+    for name, connection in imported_connections_to_save.items():
+        cli_console.step(f"Saving [{name}] connection in Snowflake CLI's config.")
+        add_connection_to_proper_file(name, ConnectionConfig.from_dict(connection))
+
+    if default_cli_connection_name in imported_connections_to_save:
+        cli_console.step(
+            f"Setting [{default_cli_connection_name}] connection as Snowflake CLI's default connection."
+        )
+        set_config_value(
+            section=None,
+            key="default_connection_name",
+            value=default_cli_connection_name,
+        )