snowflake-cli 3.5.0__py3-none-any.whl → 3.6.0__py3-none-any.whl

snowflake/cli/__about__.py

@@ -14,4 +14,16 @@
 
 from __future__ import annotations
 
-VERSION = "3.5.0"
+from enum import Enum, unique
+
+VERSION = "3.6.0"
+
+
+@unique
+class CLIInstallationSource(Enum):
+    BINARY = "binary"
+    PYPI = "pypi"
+
+
+# This variable is changed in binary release script
+INSTALLATION_SOURCE = CLIInstallationSource.PYPI

snowflake/cli/_app/commands_registration/builtin_plugins.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from snowflake.cli._plugins.auth.keypair import plugin_spec as auth_plugin_spec
 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
@@ -33,6 +34,7 @@ from snowflake.cli._plugins.workspace import plugin_spec as workspace_plugin_spe
 # plugin name to plugin spec
 def get_builtin_plugin_name_to_plugin_spec():
     plugin_specs = {
+        "auth": auth_plugin_spec,
         "connection": connection_plugin_spec,
         "helpers": migrate_plugin_spec,
         "spcs": spcs_plugin_spec,

snowflake/cli/_app/snow_connector.py

@@ -21,12 +21,11 @@ from typing import Dict, Optional
 
 import snowflake.connector
 from click.exceptions import ClickException
-from snowflake.cli.__about__ import VERSION
+from snowflake.cli import __about__
 from snowflake.cli._app.constants import (
     INTERNAL_APPLICATION_NAME,
     PARAM_APPLICATION_NAME,
 )
-from snowflake.cli._app.secret import SecretType
 from snowflake.cli._app.telemetry import command_info
 from snowflake.cli.api.config import (
     get_connection_dict,
@@ -38,6 +37,7 @@ from snowflake.cli.api.exceptions import (
     SnowflakeConnectionError,
 )
 from snowflake.cli.api.feature_flags import FeatureFlag
+from snowflake.cli.api.secret import SecretType
 from snowflake.cli.api.secure_path import SecurePath
 from snowflake.connector import SnowflakeConnection
 from snowflake.connector.errors import DatabaseError, ForbiddenError
@@ -56,6 +56,7 @@ SUPPORTED_ENV_OVERRIDES = [
     "authenticator",
     "private_key_file",
     "private_key_path",
+    "private_key_raw",
     "database",
     "schema",
     "role",
@@ -247,7 +248,7 @@ def _update_internal_application_info(connection_parameters: Dict):
     """Update internal application data if ENABLE_SEPARATE_AUTHENTICATION_POLICY_ID is enabled."""
     if FeatureFlag.ENABLE_SEPARATE_AUTHENTICATION_POLICY_ID.is_enabled():
         connection_parameters["internal_application_name"] = INTERNAL_APPLICATION_NAME
-        connection_parameters["internal_application_version"] = VERSION
+        connection_parameters["internal_application_version"] = __about__.VERSION
 
 
 def _load_pem_from_file(private_key_file: str) -> SecretType:
@@ -273,7 +274,7 @@ def _load_pem_to_der(private_key_pem: SecretType) -> SecretType:
         and private_key_passphrase.value is None
     ):
         raise ClickException(
-            "Encrypted private key, you must provide the"
+            "Encrypted private key, you must provide the "
             "passphrase in the environment variable PRIVATE_KEY_PASSPHRASE"
         )
 

snowflake/cli/_app/telemetry.py

@@ -22,7 +22,7 @@ from typing import Any, Dict, Union
 
 import click
 import typer
-from snowflake.cli.__about__ import VERSION
+from snowflake.cli import __about__
 from snowflake.cli._app.constants import PARAM_APPLICATION_NAME
 from snowflake.cli.api.cli_global_context import (
     _CliGlobalContextAccess,
@@ -40,12 +40,6 @@ 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
@@ -172,12 +166,6 @@ def _get_definition_version() -> str | None:
     return None
 
 
-def _get_installation_source() -> CLIInstallationSource:
-    if getattr(sys, "frozen", False) and hasattr(sys, "_MEIPASS"):
-        return CLIInstallationSource.BINARY
-    return CLIInstallationSource.PYPI
-
-
 def _get_ci_environment_type() -> str:
     if "GITHUB_ACTIONS" in os.environ:
         return "GITHUB_ACTIONS"
@@ -214,8 +202,8 @@ class CLITelemetryClient:
     ) -> Dict[str, Any]:
         data = {
             CLITelemetryField.SOURCE: PARAM_APPLICATION_NAME,
-            CLITelemetryField.INSTALLATION_SOURCE: _get_installation_source().value,
-            CLITelemetryField.VERSION_CLI: VERSION,
+            CLITelemetryField.INSTALLATION_SOURCE: __about__.INSTALLATION_SOURCE.value,
+            CLITelemetryField.VERSION_CLI: __about__.VERSION,
             CLITelemetryField.VERSION_OS: platform.platform(),
             CLITelemetryField.VERSION_PYTHON: python_version(),
             CLITelemetryField.COMMAND_CI_ENVIRONMENT: _get_ci_environment_type(),

snowflake/cli/_app/version_check.py

@@ -8,6 +8,8 @@ from snowflake.cli.api.console import cli_console
 from snowflake.cli.api.secure_path import SecurePath
 from snowflake.connector.config_manager import CONFIG_MANAGER
 
+REPOSITORY_URL = "https://pypi.org/pypi/snowflake-cli/json"
+
 
 def get_new_version_msg() -> str | None:
     last = _VersionCache().get_last_version()
@@ -45,9 +47,7 @@ class _VersionCache:
     @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 = requests.get(REPOSITORY_URL, headers=headers, timeout=3)
         response.raise_for_status()
         return response.json()["info"]["version"]
 
@@ -60,7 +60,7 @@ class _VersionCache:
 
     def _read_latest_version(self) -> Version | None:
         if self._cache_file.exists():
-            data = json.loads(self._cache_file.read_text())
+            data = json.loads(self._cache_file.read_text(file_size_limit_mb=1))
             now = time.time()
             if data[_VersionCache._last_time] > now - 60 * 60:
                 return Version(data[_VersionCache._version])

snowflake/cli/_plugins/auth/__init__.py

@@ -0,0 +1,11 @@
+from snowflake.cli._plugins.auth.keypair.commands import app as keypair_app
+from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
+from snowflake.cli.api.feature_flags import FeatureFlag
+
+app = SnowTyperFactory(
+    name="auth",
+    help="Manages authentication methods.",
+    is_hidden=lambda: FeatureFlag.ENABLE_AUTH_KEYPAIR.is_disabled(),
+)
+
+app.add_typer(keypair_app)

snowflake/cli/_plugins/auth/keypair/commands.py

@@ -0,0 +1,151 @@
+from pathlib import Path
+
+import typer
+from snowflake.cli._plugins.auth.keypair.manager import AuthManager, PublicKeyProperty
+from snowflake.cli.api.commands.flags import SecretTypeParser
+from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
+from snowflake.cli.api.output.types import (
+    CollectionResult,
+    CommandResult,
+    MessageResult,
+    SingleQueryResult,
+)
+from snowflake.cli.api.secret import SecretType
+from snowflake.cli.api.secure_path import SecurePath
+
+app = SnowTyperFactory(
+    name="keypair",
+    help="Manages authentication.",
+)
+
+
+KEY_PAIR_DEFAULT_PATH = "~/.ssh"
+
+
+def _show_connection_name_prompt(ctx: typer.Context, value: str):
+    for param in ctx.command.params:
+        if param.name == "connection_name":
+            if value:
+                param.prompt = "Enter connection name"
+            break
+    return value
+
+
+_new_connection_option = typer.Option(
+    True,
+    help="Create a new connection.",
+    prompt="Create a new connection?",
+    callback=_show_connection_name_prompt,
+    show_default=False,
+    hidden=True,
+)
+
+
+_connection_name_option = typer.Option(
+    None,
+    help="The new connection name.",
+    prompt=False,
+    show_default=False,
+    hidden=True,
+)
+
+
+_key_length_option = typer.Option(
+    2048,
+    "--key-length",
+    help="The RSA key length.",
+    prompt="Enter key length",
+)
+
+
+_output_path_option = typer.Option(
+    KEY_PAIR_DEFAULT_PATH,
+    "--output-path",
+    help="The output path for private and public keys",
+    prompt="Enter output path",
+)
+
+
+_private_key_passphrase_option = typer.Option(
+    "",
+    "--private-key-passphrase",
+    help="The private key passphrase.",
+    click_type=SecretTypeParser(),
+    prompt="Enter private key passphrase",
+    hide_input=True,
+    show_default=False,
+)
+
+
+@app.command("setup", requires_connection=True)
+def setup(
+    new_connection: bool = _new_connection_option,
+    connection_name: str = _connection_name_option,
+    key_length: int = _key_length_option,
+    output_path: Path = _output_path_option,
+    private_key_passphrase: SecretType = _private_key_passphrase_option,
+    **options,
+):
+    """
+    Generates the key pair, sets the public key for the user in Snowflake, and creates or updates the connection.
+    """
+    AuthManager().setup(
+        connection_name=connection_name,
+        key_length=key_length,
+        output_path=SecurePath(output_path),
+        private_key_passphrase=private_key_passphrase,
+    )
+    return MessageResult(f"Setup completed.")
+
+
+@app.command("rotate", requires_connection=True)
+def rotate(
+    key_length: int = _key_length_option,
+    output_path: Path = _output_path_option,
+    private_key_passphrase: SecretType = _private_key_passphrase_option,
+    **options,
+):
+    """
+    Rotates the key for the connection. Generates the key pair, sets the public key for the user in Snowflake, and creates or updates the connection.
+    """
+    AuthManager().rotate(
+        key_length=key_length,
+        output_path=SecurePath(output_path),
+        private_key_passphrase=private_key_passphrase,
+    )
+    return MessageResult(f"Rotate completed.")
+
+
+@app.command("list", requires_connection=True)
+def list_keys(**options) -> CommandResult:
+    """
+    Lists the public keys set for the user.
+    """
+    return CollectionResult(AuthManager().list_keys())
+
+
+@app.command("remove", requires_connection=True)
+def remove(
+    public_key_property: PublicKeyProperty = typer.Option(
+        ...,
+        "--key-id",
+        help=f"Local path to the template directory or a URL to Git repository with templates.",
+        show_default=False,
+    ),
+    **options,
+):
+    """
+    Removes the public key for the user.
+    """
+    return SingleQueryResult(
+        AuthManager().remove_public_key(public_key_property=public_key_property)
+    )
+
+
+@app.command("status", requires_connection=True)
+def status(**options):
+    """
+    Verifies the key pair configuration and tests the connection.
+    """
+    AuthManager().status()
+    return MessageResult("Status check completed.")

snowflake/cli/_plugins/auth/keypair/manager.py

@@ -0,0 +1,331 @@
+from enum import Enum
+from typing import Dict, List, Optional, Tuple
+
+from click import ClickException
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+from snowflake.cli._plugins.object.manager import ObjectManager
+from snowflake.cli.api import exceptions
+from snowflake.cli.api.cli_global_context import (
+    _CliGlobalContextAccess,
+    get_cli_context,
+)
+from snowflake.cli.api.config import (
+    connection_exists,
+    get_connection_dict,
+    set_config_value,
+)
+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.secret import SecretType
+from snowflake.cli.api.secure_path import SecurePath
+from snowflake.cli.api.sql_execution import SqlExecutionMixin
+from snowflake.connector import DictCursor
+from snowflake.connector.cursor import SnowflakeCursor
+
+
+class PublicKeyProperty(Enum):
+    RSA_PUBLIC_KEY = "RSA_PUBLIC_KEY"
+    RSA_PUBLIC_KEY_2 = "RSA_PUBLIC_KEY_2"
+
+
+class AuthManager(SqlExecutionMixin):
+    def setup(
+        self,
+        connection_name: Optional[str],
+        key_length: int,
+        output_path: SecurePath,
+        private_key_passphrase: SecretType,
+    ):
+        # When the user provide new connection name
+        if connection_name and connection_exists(connection_name):
+            raise ClickException(
+                f"Connection with name {connection_name} already exists."
+            )
+
+        cli_context = get_cli_context()
+        # When the use not provide connection name, so we overwrite the current connection
+        if not connection_name:
+            connection_name = cli_context.connection_context.connection_name
+
+        key_name = AuthManager._get_free_key_name(output_path, connection_name)  # type: ignore[arg-type]
+        self._generate_key_pair_and_set_public_key(
+            user=cli_context.connection.user,
+            key_length=key_length,
+            output_path=output_path,
+            key_name=key_name,  # type: ignore[arg-type]
+            private_key_passphrase=private_key_passphrase,
+        )
+
+        self._create_or_update_connection(
+            current_connection=cli_context.connection_context.connection_name,
+            connection_name=connection_name,  # type: ignore[arg-type]
+            private_key_path=self._get_private_key_path(
+                output_path=output_path, key_name=key_name  # type: ignore[arg-type]
+            ),
+        )
+
+    def rotate(
+        self,
+        key_length: int,
+        output_path: SecurePath,
+        private_key_passphrase: SecretType,
+    ):
+        cli_context = get_cli_context()
+        connection_name = cli_context.connection_context.connection_name
+
+        self._ensure_connection_has_private_key(
+            cli_context.connection_context.connection_name
+        )
+
+        public_key, public_key_2 = self._get_public_keys()
+
+        if not public_key and not public_key_2:
+            raise ClickException("No public key found. Use the setup command first.")
+
+        if public_key_2:
+            self.set_public_key(
+                cli_context.connection.user,
+                PublicKeyProperty.RSA_PUBLIC_KEY,
+                public_key_2,
+            )
+
+        key_name = AuthManager._get_free_key_name(output_path, connection_name)
+        public_key = self._generate_keys_and_return_public_key(
+            key_length=key_length,
+            output_path=output_path,
+            key_name=key_name,
+            private_key_passphrase=private_key_passphrase,
+        )
+        self.set_public_key(
+            cli_context.connection.user, PublicKeyProperty.RSA_PUBLIC_KEY_2, public_key
+        )
+        self._create_or_update_connection(
+            current_connection=cli_context.connection_context.connection_name,
+            connection_name=connection_name,
+            private_key_path=self._get_private_key_path(
+                output_path=output_path, key_name=key_name
+            ),
+        )
+
+    def _generate_key_pair_and_set_public_key(
+        self,
+        user: str,
+        key_length: int,
+        output_path: SecurePath,
+        key_name: str,
+        private_key_passphrase: SecretType,
+    ):
+        public_key_exists, public_key_2_exists = self._get_public_keys()
+
+        if public_key_exists or public_key_2_exists:
+            raise exceptions.CouldNotSetKeyPairError()
+
+        if not output_path.exists():
+            output_path.mkdir(parents=True)
+
+        public_key = self._generate_keys_and_return_public_key(
+            key_length=key_length,
+            output_path=output_path,
+            key_name=key_name,  # type: ignore[arg-type]
+            private_key_passphrase=private_key_passphrase,
+        )
+        self.set_public_key(user, PublicKeyProperty.RSA_PUBLIC_KEY, public_key)
+
+    def list_keys(self) -> List[Dict]:
+        key_properties = [
+            "RSA_PUBLIC_KEY",
+            "RSA_PUBLIC_KEY_FP",
+            "RSA_PUBLIC_KEY_LAST_SET_TIME",
+            "RSA_PUBLIC_KEY_2",
+            "RSA_PUBLIC_KEY_2_FP",
+            "RSA_PUBLIC_KEY_2_LAST_SET_TIME",
+        ]
+        cursor = ObjectManager(connection=self._conn).describe(
+            object_type=ObjectType.USER.value.sf_name,
+            fqn=FQN.from_string(self._conn.user),
+            cursor_class=DictCursor,
+        )
+        only_public_key_properties = [
+            p for p in cursor.fetchall() if p.get("property") in key_properties
+        ]
+        return only_public_key_properties
+
+    def set_public_key(
+        self, user: str, public_key_property: PublicKeyProperty, public_key: str
+    ) -> SnowflakeCursor:
+        return self.execute_query(
+            f"ALTER USER {user} SET {public_key_property.value}='{public_key}'"
+        )
+
+    def remove_public_key(
+        self, public_key_property: PublicKeyProperty
+    ) -> SnowflakeCursor:
+        cli_context = get_cli_context()
+        return self.execute_query(
+            f"ALTER USER {cli_context.connection.user} UNSET {public_key_property.value}"
+        )
+
+    def status(self):
+        cli_context = get_cli_context()
+        self._ensure_connection_has_private_key(
+            cli_context.connection_context.connection_name
+        )
+        cli_console.step("Private key set for connection - OK")
+        self._check_connection(cli_context)
+        cli_console.step("Test connection - OK")
+
+    def extend_connection_add(
+        self,
+        connection_name: str,
+        connection_options: Dict,
+        key_length: int,
+        output_path: SecurePath,
+        private_key_passphrase: SecretType,
+    ) -> Dict:
+        key_name = AuthManager._get_free_key_name(output_path, connection_name)
+
+        self._generate_key_pair_and_set_public_key(
+            user=connection_options["user"],
+            key_length=key_length,
+            output_path=output_path,
+            key_name=key_name,
+            private_key_passphrase=private_key_passphrase,
+        )
+
+        connection_options["authenticator"] = "SNOWFLAKE_JWT"
+        connection_options["private_key_file"] = str(
+            self._get_private_key_path(output_path=output_path, key_name=key_name).path
+        )
+        if connection_options.get("password"):
+            del connection_options["password"]
+
+        return connection_options
+
+    @staticmethod
+    def _ensure_connection_has_private_key(connection_name: str) -> None:
+        connection = get_connection_dict(connection_name)
+        if not connection.get("private_key_file") and not connection.get(
+            "private_key_path"
+        ):
+            raise ClickException(
+                f"The private key is not set in {connection_name} connection."
+            )
+
+    @staticmethod
+    def _check_connection(cli_context: _CliGlobalContextAccess) -> None:
+        cli_context.connection
+
+    def _get_public_keys(self) -> Tuple[str, str]:
+        keys = self.list_keys()
+        public_key = ""
+        public_key_2 = ""
+        for p in keys:
+            if (
+                p.get("property") == PublicKeyProperty.RSA_PUBLIC_KEY.value
+                and p.get("value") != "null"
+            ):
+                public_key = p.get("value")  # type: ignore
+            if (
+                p.get("property") == PublicKeyProperty.RSA_PUBLIC_KEY_2.value
+                and p.get("value") != "null"
+            ):
+                public_key_2 = p.get("value")  # type: ignore
+        return public_key, public_key_2
+
+    @staticmethod
+    def _generate_keys_and_return_public_key(
+        key_length: int,
+        output_path: SecurePath,
+        key_name: str,
+        private_key_passphrase: SecretType,
+    ) -> str:
+        private_key = rsa.generate_private_key(
+            public_exponent=65537,
+            key_size=key_length,
+        )
+
+        if private_key_passphrase:
+            pem = SecretType(
+                private_key.private_bytes(
+                    encoding=serialization.Encoding.PEM,
+                    format=serialization.PrivateFormat.PKCS8,
+                    encryption_algorithm=serialization.BestAvailableEncryption(
+                        private_key_passphrase.value.encode("utf-8")
+                    ),
+                )
+            )
+            cli_console.message(
+                "Set the `PRIVATE_KEY_PASSPHRASE` environment variable before using the connection."
+            )
+        else:
+            pem = SecretType(
+                private_key.private_bytes(
+                    encoding=serialization.Encoding.PEM,
+                    format=serialization.PrivateFormat.PKCS8,
+                    encryption_algorithm=serialization.NoEncryption(),
+                )
+            )
+
+        with AuthManager._get_private_key_path(output_path, key_name).open(
+            mode="wb"
+        ) as file:
+            file.write(pem.value)
+
+        public_key = private_key.public_key()
+        public_pem = public_key.public_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PublicFormat.SubjectPublicKeyInfo,
+        )
+        with AuthManager._get_public_key_path(output_path, key_name).open(
+            mode="wb"
+        ) as file:
+            file.write(public_pem)
+
+        return public_pem.decode("utf-8")
+
+    @staticmethod
+    def _get_free_key_name(output_path: SecurePath, key_name: str) -> str:
+        new_private_key = f"{key_name}.p8"
+        new_public_key = f"{key_name}.pub"
+        new_key_name = key_name
+        counter = 1
+
+        while (
+            (output_path / new_private_key).exists()
+            and (output_path / new_public_key).exists()
+            and counter <= 100
+        ):
+            new_key_name = f"{key_name}_{counter}"
+            new_private_key = f"{new_key_name}.p8"
+            new_public_key = f"{new_key_name}.pub"
+            counter += 1
+
+        if counter == 100:
+            raise ClickException(
+                "Too many key pairs with the same name in the output directory."
+            )
+
+        return new_key_name
+
+    @staticmethod
+    def _get_private_key_path(output_path: SecurePath, key_name: str) -> SecurePath:
+        return (output_path / f"{key_name}.p8").resolve()
+
+    @staticmethod
+    def _get_public_key_path(output_path: SecurePath, key_name: str) -> SecurePath:
+        return (output_path / f"{key_name}.pub").resolve()
+
+    @staticmethod
+    def _create_or_update_connection(
+        current_connection: Optional[str],
+        connection_name: str,
+        private_key_path: SecurePath,
+    ):
+        connection = get_connection_dict(current_connection)
+        connection.pop("password", None)
+        connection["authenticator"] = "SNOWFLAKE_JWT"
+        connection["private_key_file"] = str(private_key_path.path)
+
+        set_config_value(["connections", connection_name], value=connection)

snowflake/cli/_plugins/auth/keypair/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.auth import app
+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=app.create_instance(),
+    )