snowflake-cli-labs 3.0.0rc5__py3-none-any.whl → 3.0.2__py3-none-any.whl

snowflake/cli/_plugins/connection/commands.py

@@ -1,353 +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 logging
-import os.path
-
-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._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_to_proper_file,
-    connection_exists,
-    get_all_connections,
-    get_connection_dict,
-    get_default_connection_name,
-    set_config_value,
-)
-from snowflake.cli.api.console import cli_console
-from snowflake.cli.api.constants import ObjectType
-from snowflake.cli.api.output.types import (
-    CollectionResult,
-    CommandResult,
-    MessageResult,
-    ObjectResult,
-)
-from snowflake.connector import ProgrammingError
-from snowflake.connector.constants import CONNECTIONS_FILE
-
-app = SnowTyperFactory(
-    name="connection",
-    help="Manages connections to Snowflake.",
-)
-log = logging.getLogger(__name__)
-
-
-class EmptyInput:
-    def __repr__(self):
-        return "optional"
-
-
-class OptionalPrompt(StringParamType):
-    def convert(self, value, param, ctx):
-        return None if isinstance(value, EmptyInput) else value
-
-
-def _mask_password(connection_params: dict):
-    if "password" in connection_params:
-        connection_params["password"] = "****"
-    return connection_params
-
-
-@app.command(name="list")
-def list_connections(**options) -> CommandResult:
-    """
-    Lists configured connections.
-    """
-    connections = get_all_connections()
-    default_connection = get_default_connection_name()
-    result = (
-        {
-            "connection_name": connection_name,
-            "parameters": _mask_password(
-                connection_config.to_dict_of_known_non_empty_values()
-            ),
-            "is_default": connection_name == default_connection,
-        }
-        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)
-
-
-def require_integer(field_name: str):
-    def callback(value: str):
-        if value is None:
-            return None
-        if value.strip().isdigit():
-            return value.strip()
-        raise ClickException(f"Value of {field_name} must be integer")
-
-    return callback
-
-
-def _password_callback(ctx: Context, param: Parameter, value: str):
-    if value and ctx.get_parameter_source(param.name) == ParameterSource.COMMANDLINE:  # type: ignore
-        cli_console.warning(PLAIN_PASSWORD_MSG)
-
-    return value
-
-
-@app.command()
-def add(
-    connection_name: str = typer.Option(
-        None,
-        "--connection-name",
-        "-n",
-        prompt="Name for this connection",
-        help="Name of the new connection.",
-        show_default=False,
-        callback=strip_if_value_present,
-    ),
-    account: str = typer.Option(
-        None,
-        "--account",
-        "-a",
-        "--accountname",
-        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,
-        "--user",
-        "-u",
-        "--username",
-        prompt="Snowflake username",
-        show_default=False,
-        help="Username to connect to Snowflake.",
-        callback=strip_if_value_present,
-    ),
-    password: str = typer.Option(
-        EmptyInput(),
-        "--password",
-        "-p",
-        click_type=OptionalPrompt(),
-        callback=_password_callback,
-        prompt="Snowflake password",
-        help="Snowflake password.",
-        hide_input=True,
-    ),
-    role: str = typer.Option(
-        EmptyInput(),
-        "--role",
-        "-r",
-        click_type=OptionalPrompt(),
-        prompt="Role for the connection",
-        help="Role to use on Snowflake.",
-        callback=strip_if_value_present,
-    ),
-    warehouse: str = typer.Option(
-        EmptyInput(),
-        "--warehouse",
-        "-w",
-        click_type=OptionalPrompt(),
-        prompt="Warehouse for the connection",
-        help="Warehouse to use on Snowflake.",
-        callback=strip_if_value_present,
-    ),
-    database: str = typer.Option(
-        EmptyInput(),
-        "--database",
-        "-d",
-        click_type=OptionalPrompt(),
-        prompt="Database for the connection",
-        help="Database to use on Snowflake.",
-        callback=strip_if_value_present,
-    ),
-    schema: str = typer.Option(
-        EmptyInput(),
-        "--schema",
-        "-s",
-        click_type=OptionalPrompt(),
-        prompt="Schema for the connection",
-        help="Schema to use on Snowflake.",
-        callback=strip_if_value_present,
-    ),
-    host: str = typer.Option(
-        EmptyInput(),
-        "--host",
-        "-h",
-        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(),
-        "--port",
-        "-P",
-        click_type=OptionalPrompt(),
-        prompt="Connection port",
-        help="Port to communicate with on the host.",
-        callback=require_integer(field_name="port"),
-    ),
-    region: str = typer.Option(
-        EmptyInput(),
-        "--region",
-        "-R",
-        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(),
-        "--authenticator",
-        "-A",
-        click_type=OptionalPrompt(),
-        prompt="Authentication method",
-        help="Chosen authenticator, if other than password-based",
-    ),
-    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(),
-        "--token-file-path",
-        "-t",
-        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,
-        "--default",
-        is_flag=True,
-        help="If provided the connection will be configured as default connection.",
-    ),
-    **options,
-) -> CommandResult:
-    """Adds a connection to configuration file."""
-    if connection_exists(connection_name):
-        raise ClickException(f"Connection {connection_name} already exists")
-
-    connections_file = add_connection_to_proper_file(
-        connection_name,
-        ConnectionConfig(
-            account=account,
-            user=user,
-            password=password,
-            host=host,
-            region=region,
-            port=port,
-            database=database,
-            schema=schema,
-            warehouse=warehouse,
-            role=role,
-            authenticator=authenticator,
-            private_key_file=private_key_file,
-            token_file_path=token_file_path,
-        ),
-    )
-    if set_as_default:
-        set_config_value(
-            section=None, key="default_connection_name", value=connection_name
-        )
-
-    return MessageResult(
-        f"Wrote new connection {connection_name} to {connections_file}"
-    )
-
-
-@app.command(requires_connection=True)
-def test(
-    **options,
-) -> CommandResult:
-    """
-    Tests the connection to Snowflake.
-    """
-
-    # Test connection
-    cli_context = get_cli_context()
-    conn = cli_context.connection
-
-    # Test session attributes
-    om = ObjectManager()
-    try:
-        # "use database" operation changes schema to default "public",
-        # so to test schema set up by user we need to copy it here:
-        schema = conn.schema
-
-        if conn.role:
-            om.use(object_type=ObjectType.ROLE, name=f'"{conn.role}"')
-        if conn.database:
-            om.use(object_type=ObjectType.DATABASE, name=f'"{conn.database}"')
-        if schema:
-            om.use(object_type=ObjectType.SCHEMA, name=f'"{schema}"')
-        if conn.warehouse:
-            om.use(object_type=ObjectType.WAREHOUSE, name=f'"{conn.warehouse}"')
-
-    except ProgrammingError as err:
-        raise ClickException(str(err))
-
-    conn_ctx = cli_context.connection_context
-    result = {
-        "Connection name": conn_ctx.connection_name,
-        "Status": "OK",
-        "Host": conn.host,
-        "Account": conn.account,
-        "User": conn.user,
-        "Role": f'{conn.role or "not set"}',
-        "Database": f'{conn.database or "not set"}',
-        "Warehouse": f'{conn.warehouse or "not set"}',
-    }
-
-    if conn_ctx.enable_diag:
-        result["Diag Report Location"] = os.path.join(
-            conn_ctx.diag_log_path, "SnowflakeConnectionTestReport.txt"
-        )
-
-    return ObjectResult(result)
-
-
-@app.command(requires_connection=False)
-def set_default(
-    name: str = typer.Argument(
-        help="Name of the connection, as defined in your `config.toml`",
-        show_default=False,
-    ),
-    **options,
-):
-    """Changes default connection to provided value."""
-    get_connection_dict(connection_name=name)
-    set_config_value(section=None, key="default_connection_name", value=name)
-    return MessageResult(f"Default connection set to: {name}")

snowflake/cli/_plugins/connection/plugin_spec.py

@@ -1,30 +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 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/connection/util.py

@@ -1,195 +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 json
-import logging
-import os
-from typing import Optional
-
-from click.exceptions import ClickException
-from snowflake.connector import SnowflakeConnection
-from snowflake.connector.cursor import DictCursor
-
-log = logging.getLogger(__name__)
-
-REGIONLESS_QUERY = """
-    select value['value'] as REGIONLESS from table(flatten(
-        input => parse_json(SYSTEM$BOOTSTRAP_DATA_REQUEST()),
-        path => 'clientParamsInfo'
-    )) 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 MissingConnectionAccountError(ClickException):
-    def __init__(self, conn: SnowflakeConnection):
-        super().__init__(
-            "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)"
-        )
-
-
-def is_regionless_redirect(conn: SnowflakeConnection) -> bool:
-    """
-    Determines if the deployment this connection refers to uses
-    regionless URLs in Snowsight (/orgname/account) or regional URLs
-    (/region/account). If we cannot determine the correct value we
-    assume it's regionless, as this is true for most production deployments.
-    """
-    try:
-        *_, cursor = conn.execute_string(REGIONLESS_QUERY, cursor_class=DictCursor)
-        return cursor.fetchone()["REGIONLESS"].lower() == "true"
-    except:
-        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.
-    This could be a region or it could be an organization, depending
-    on whether or not the underlying deployment uses regionless URLs.
-    """
-    if is_regionless_redirect(conn):
-        *_, cursor = conn.execute_string(
-            f"select system$return_current_org_name()", cursor_class=DictCursor
-        )
-        return cursor.fetchone()["SYSTEM$RETURN_CURRENT_ORG_NAME()"]
-
-    return get_region(conn)
-
-
-def get_account(conn: SnowflakeConnection) -> str:
-    """
-    Determines the account that this connection refers to.
-    """
-    try:
-        *_, cursor = conn.execute_string(
-            f"select current_account_name()", cursor_class=DictCursor
-        )
-        return cursor.fetchone()["CURRENT_ACCOUNT_NAME()"].lower()
-    except Exception as e:
-        # try to extract the account from the connection information
-        if conn.account:
-            return conn.account
-
-        if conn.host:
-            host_parts = conn.host.split(".")
-            return host_parts[0]
-
-        raise MissingConnectionAccountError(conn)
-
-
-def get_snowsight_host(conn: SnowflakeConnection) -> str:
-    try:
-        *_, cursor = conn.execute_string(
-            f"select system$get_snowsight_host()", cursor_class=DictCursor
-        )
-        return cursor.fetchone()["SYSTEM$GET_SNOWSIGHT_HOST()"]
-    except Exception as e:
-        # if we cannot determine the host, assume we're on prod
-        return "https://app.snowflake.com"
-
-
-def make_snowsight_url(conn: SnowflakeConnection, path: str) -> str:
-    """
-    Returns a URL on the correct Snowsight instance for the connected account.
-    The path that is passed in must already be properly URL-encoded, and
-    can optionally contain a hash/fragment (e.g. #).
-
-    See also identifier_for_url.
-    """
-    snowsight_host = get_snowsight_host(conn)
-    deployment = get_context(conn)
-    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/cortex/__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.