snowflake-cli 2.8.2__py3-none-any.whl

snowflake/cli/plugins/connection/commands.py

@@ -0,0 +1,330 @@
+# 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.api.cli_global_context import 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,
+    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.cli.plugins.object.manager import ObjectManager
+from snowflake.connector import ProgrammingError
+from snowflake.connector.config_manager import CONFIG_MANAGER
+
+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()
+    )
+    return CollectionResult(result)
+
+
+def require_integer(field_name: str):
+    def callback(value: str):
+        if value is None:
+            return None
+        if value.isdigit():
+            return value
+        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,
+    ),
+    account: str = typer.Option(
+        None,
+        "--account",
+        "-a",
+        "--accountname",
+        prompt="Snowflake account name",
+        help="Account name to use when authenticating with Snowflake.",
+        show_default=False,
+    ),
+    user: str = typer.Option(
+        None,
+        "--user",
+        "-u",
+        "--username",
+        prompt="Snowflake username",
+        show_default=False,
+        help="Username to connect to Snowflake.",
+    ),
+    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.",
+    ),
+    warehouse: str = typer.Option(
+        EmptyInput(),
+        "--warehouse",
+        "-w",
+        click_type=OptionalPrompt(),
+        prompt="Warehouse for the connection",
+        help="Warehouse to use on Snowflake.",
+    ),
+    database: str = typer.Option(
+        EmptyInput(),
+        "--database",
+        "-d",
+        click_type=OptionalPrompt(),
+        prompt="Database for the connection",
+        help="Database to use on Snowflake.",
+    ),
+    schema: str = typer.Option(
+        EmptyInput(),
+        "--schema",
+        "-s",
+        click_type=OptionalPrompt(),
+        prompt="Schema for the connection",
+        help="Schema to use on Snowflake.",
+    ),
+    host: str = typer.Option(
+        EmptyInput(),
+        "--host",
+        "-h",
+        click_type=OptionalPrompt(),
+        prompt="Connection host",
+        help="Host name the connection attempts to connect to Snowflake.",
+    ),
+    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.",
+    ),
+    authenticator: str = typer.Option(
+        EmptyInput(),
+        "--authenticator",
+        "-A",
+        click_type=OptionalPrompt(),
+        prompt="Authentication method",
+        help="Chosen authenticator, if other than password-based",
+    ),
+    private_key_path: str = typer.Option(
+        EmptyInput(),
+        "--private-key",
+        "-k",
+        click_type=OptionalPrompt(),
+        prompt="Path to private key file",
+        help="Path to file containing private key",
+    ),
+    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",
+    ),
+    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")
+
+    add_connection(
+        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_path=private_key_path,
+            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 {CONFIG_MANAGER.file_path}"
+    )
+
+
+@app.command(requires_connection=True)
+def test(
+    **options,
+) -> CommandResult:
+    """
+    Tests the connection to Snowflake.
+    """
+
+    # Test connection
+    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`"
+    ),
+    **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

@@ -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.api.plugins.command import (
+    SNOWCLI_ROOT_COMMAND_PATH,
+    CommandSpec,
+    CommandType,
+    plugin_hook_impl,
+)
+from snowflake.cli.plugins.connection import commands
+
+
+@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

@@ -0,0 +1,179 @@
+# 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
+
+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}"

snowflake/cli/plugins/cortex/__init__.py

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