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

snowflake/cli/app/snow_connector.py

@@ -1,243 +0,0 @@
-# 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
-
-import contextlib
-import logging
-import os
-from typing import Dict, Optional
-
-import snowflake.connector
-from click.exceptions import ClickException
-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,
-    SnowflakeConnectionError,
-)
-from snowflake.cli.api.secure_path import SecurePath
-from snowflake.cli.app.constants import (
-    PARAM_APPLICATION_NAME,
-)
-from snowflake.cli.app.telemetry import command_info
-from snowflake.connector import SnowflakeConnection
-from snowflake.connector.errors import DatabaseError, ForbiddenError
-
-log = logging.getLogger(__name__)
-
-
-ENCRYPTED_PKCS8_PK_HEADER = b"-----BEGIN ENCRYPTED PRIVATE KEY-----"
-UNENCRYPTED_PKCS8_PK_HEADER = b"-----BEGIN PRIVATE KEY-----"
-
-
-def connect_to_snowflake(
-    temporary_connection: bool = False,
-    mfa_passcode: Optional[str] = None,
-    enable_diag: Optional[bool] = False,
-    diag_log_path: Optional[str] = None,
-    diag_allowlist_path: Optional[str] = None,
-    connection_name: Optional[str] = None,
-    **overrides,
-) -> SnowflakeConnection:
-    if temporary_connection and connection_name:
-        raise ClickException("Can't use connection name and temporary connection.")
-
-    using_session_token = (
-        "session_token" in overrides and overrides["session_token"] is not None
-    )
-    using_master_token = (
-        "master_token" in overrides and overrides["master_token"] is not None
-    )
-    _raise_errors_related_to_session_token(
-        temporary_connection, using_session_token, using_master_token
-    )
-
-    if connection_name:
-        connection_parameters = get_connection_dict(connection_name)
-    elif temporary_connection:
-        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():
-        # Command line override case
-        if value:
-            connection_parameters[key] = value
-            continue
-
-        # Generic environment variable case, apply only if value not passed via flag or connection variable
-        generic_env_value = os.environ.get(f"SNOWFLAKE_{key}".upper())
-        if key not in connection_parameters and generic_env_value:
-            connection_parameters[key] = generic_env_value
-            continue
-
-    # Clean up connection params
-    connection_parameters = {
-        k: v for k, v in connection_parameters.items() if v is not None
-    }
-
-    connection_parameters = update_connection_details_with_private_key(
-        connection_parameters
-    )
-
-    if mfa_passcode:
-        connection_parameters["passcode"] = mfa_passcode
-
-    if connection_parameters.get("authenticator") == "username_password_mfa":
-        connection_parameters["client_request_mfa_token"] = True
-
-    if enable_diag:
-        connection_parameters["enable_connection_diag"] = enable_diag
-        if diag_log_path:
-            connection_parameters["connection_diag_log_path"] = diag_log_path
-        if diag_allowlist_path:
-            connection_parameters[
-                "connection_diag_allowlist_path"
-            ] = diag_allowlist_path
-
-    # Make sure the connection is not closed if it was shared to the SnowCLI, instead of being created in the SnowCLI
-    _avoid_closing_the_connection_if_it_was_shared(
-        using_session_token, using_master_token, connection_parameters
-    )
-
-    _update_connection_application_name(connection_parameters)
-
-    try:
-        # Whatever output is generated when creating connection,
-        # we don't want it in our output. This is particularly important
-        # for cases when external browser and json format are used.
-        # Redirecting both stdout and stderr for offline usage.
-        with contextlib.redirect_stdout(None), contextlib.redirect_stderr(None):
-            return snowflake.connector.connect(
-                application=command_info(),
-                **connection_parameters,
-            )
-    except ForbiddenError as err:
-        raise SnowflakeConnectionError(err)
-    except DatabaseError as err:
-        raise InvalidConnectionConfiguration(err.msg)
-
-
-def _avoid_closing_the_connection_if_it_was_shared(
-    using_session_token: bool, using_master_token: bool, connection_parameters: Dict
-):
-    if using_session_token and using_master_token:
-        connection_parameters["server_session_keep_alive"] = True
-
-
-def _raise_errors_related_to_session_token(
-    temporary_connection: bool, using_session_token: bool, using_master_token: bool
-):
-    if not temporary_connection and (using_session_token or using_master_token):
-        raise ClickException(
-            "When using a session or master token, you must use a temporary connection"
-        )
-    if using_session_token and not using_master_token:
-        raise ClickException(
-            "When using a session token, you must provide the corresponding master token"
-        )
-    if using_master_token and not using_session_token:
-        raise ClickException(
-            "When using a master token, you must provide the corresponding session token"
-        )
-
-
-def update_connection_details_with_private_key(connection_parameters: Dict):
-    if "private_key_path" in connection_parameters:
-        if connection_parameters.get("authenticator") == "SNOWFLAKE_JWT":
-            private_key = _load_pem_to_der(connection_parameters["private_key_path"])
-            connection_parameters["private_key"] = private_key
-            del connection_parameters["private_key_path"]
-        else:
-            raise ClickException(
-                "Private Key authentication requires authenticator set to SNOWFLAKE_JWT"
-            )
-    return connection_parameters
-
-
-def _update_connection_application_name(connection_parameters: Dict):
-    """Update version and name of app handling connection."""
-    connection_application_params = {
-        "application_name": PARAM_APPLICATION_NAME,
-    }
-    connection_parameters.update(connection_application_params)
-
-
-def _load_pem_to_der(private_key_path: str) -> bytes:
-    """
-    Given a private key file path (in PEM format), decode key data into DER
-    format
-    """
-
-    with SecurePath(private_key_path).open(
-        "rb", read_file_limit_mb=DEFAULT_SIZE_LIMIT_MB
-    ) as f:
-        private_key_pem = f.read()
-
-    private_key_passphrase = os.getenv("PRIVATE_KEY_PASSPHRASE", None)
-    if (
-        private_key_pem.startswith(ENCRYPTED_PKCS8_PK_HEADER)
-        and private_key_passphrase is None
-    ):
-        raise ClickException(
-            "Encrypted private key, you must provide the"
-            "passphrase in the environment variable PRIVATE_KEY_PASSPHRASE"
-        )
-
-    if not private_key_pem.startswith(
-        ENCRYPTED_PKCS8_PK_HEADER
-    ) and not private_key_pem.startswith(UNENCRYPTED_PKCS8_PK_HEADER):
-        raise ClickException(
-            "Private key provided is not in PKCS#8 format. Please use correct format."
-        )
-
-    if private_key_pem.startswith(UNENCRYPTED_PKCS8_PK_HEADER):
-        private_key_passphrase = None
-
-    return prepare_private_key(private_key_pem, private_key_passphrase)
-
-
-def prepare_private_key(private_key_pem, private_key_passphrase=None):
-    from cryptography.hazmat.backends import default_backend
-    from cryptography.hazmat.primitives.serialization import (
-        Encoding,
-        NoEncryption,
-        PrivateFormat,
-        load_pem_private_key,
-    )
-
-    private_key = load_pem_private_key(
-        private_key_pem,
-        (
-            str.encode(private_key_passphrase)
-            if private_key_passphrase is not None
-            else private_key_passphrase
-        ),
-        default_backend(),
-    )
-    return private_key.private_bytes(
-        encoding=Encoding.DER,
-        format=PrivateFormat.PKCS8,
-        encryption_algorithm=NoEncryption(),
-    )

snowflake/cli/app/telemetry.py

@@ -1,189 +0,0 @@
-# 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
-
-import platform
-import sys
-from enum import Enum, unique
-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
-from snowflake.cli.app.constants import PARAM_APPLICATION_NAME
-from snowflake.connector.telemetry import (
-    TelemetryData,
-    TelemetryField,
-)
-from snowflake.connector.time_util import get_time_millis
-
-
-@unique
-class CLITelemetryField(Enum):
-    # Basic information
-    SOURCE = "source"
-    VERSION_CLI = "version_cli"
-    VERSION_PYTHON = "version_python"
-    VERSION_OS = "version_os"
-    # Command execution context
-    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]
-
-
-def _find_command_info() -> TelemetryDict:
-    ctx = click.get_current_context()
-    command_path = ctx.command_path.split(" ")[1:]
-    return {
-        CLITelemetryField.COMMAND: command_path,
-        CLITelemetryField.COMMAND_GROUP: command_path[0],
-        CLITelemetryField.COMMAND_FLAGS: {
-            k: ctx.get_parameter_source(k).name  # type: ignore[attr-defined]
-            for k, v in ctx.params.items()
-            if v  # noqa
-        },
-        CLITelemetryField.COMMAND_OUTPUT_TYPE: ctx.params.get(
-            "format", OutputFormat.TABLE
-        ).value,
-        CLITelemetryField.PROJECT_DEFINITION_VERSION: str(_get_definition_version()),
-    }
-
-
-def _get_definition_version() -> str | None:
-    from snowflake.cli.api.cli_global_context import cli_context
-
-    if cli_context.project_definition:
-        return cli_context.project_definition.definition_version
-    return None
-
-
-def command_info() -> str:
-    info = _find_command_info()
-    command = ".".join(info[CLITelemetryField.COMMAND])
-    return f"{PARAM_APPLICATION_NAME}.{command}".upper()
-
-
-def python_version() -> str:
-    py_ver = sys.version_info
-    return f"{py_ver.major}.{py_ver.minor}.{py_ver.micro}"
-
-
-class CLITelemetryClient:
-    def __init__(self, ctx):
-        self._ctx = ctx
-
-    @staticmethod
-    def generate_telemetry_data_dict(
-        telemetry_payload: TelemetryDict,
-    ) -> Dict[str, Any]:
-        data = {
-            CLITelemetryField.SOURCE: PARAM_APPLICATION_NAME,
-            CLITelemetryField.VERSION_CLI: VERSION,
-            CLITelemetryField.VERSION_OS: platform.platform(),
-            CLITelemetryField.VERSION_PYTHON: python_version(),
-            CLITelemetryField.CONFIG_FEATURE_FLAGS: {
-                k: str(v) for k, v in get_feature_flags_section().items()
-            },
-            **_find_command_info(),
-            **telemetry_payload,
-        }
-        # To map Enum to string, so we don't have to use .value every time
-        return {getattr(k, "value", k): v for k, v in data.items()}  # type: ignore[arg-type]
-
-    @property
-    def _telemetry(self):
-        return self._ctx.connection._telemetry  # noqa
-
-    def send(self, payload: TelemetryDict):
-        if self._telemetry:
-            message = self.generate_telemetry_data_dict(payload)
-            telemetry_data = TelemetryData.from_telemetry_data_dict(
-                from_dict=message, timestamp=get_time_millis()
-            )
-            self._telemetry.try_add_log_to_batch(telemetry_data)
-
-    def flush(self):
-        self._telemetry.send_batch()
-
-
-_telemetry = CLITelemetryClient(ctx=cli_context)
-
-
-@ignore_exceptions()
-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()
-def flush_telemetry():
-    _telemetry.flush()

snowflake/cli/plugins/__init__.py

@@ -1,13 +0,0 @@
-# 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.

snowflake/cli/plugins/connection/__init__.py

@@ -1,13 +0,0 @@
-# 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.