SNOW-1980772: command for checking variables (#2127)

Author: sfc-gh-mraba

src/snowflake/cli/_plugins/helpers/commands.py

@@ -15,12 +15,14 @@
 from __future__ import annotations
 
 import logging
+import os
 from pathlib import Path
 from typing import Any, List, Optional
 
 import typer
 import yaml
 from click import ClickException
+from snowflake.cli._plugins.helpers.snowsl_vars_reader import check_env_vars
 from snowflake.cli.api.commands.snow_typer import SnowTyperFactory
 from snowflake.cli.api.config import (
     ConnectionConfig,
@@ -29,7 +31,12 @@
     set_config_value,
 )
 from snowflake.cli.api.console import cli_console
-from snowflake.cli.api.output.types import CommandResult, MessageResult
+from snowflake.cli.api.output.types import (
+    CollectionResult,
+    CommandResult,
+    MessageResult,
+    MultipleResults,
+)
 from snowflake.cli.api.project.definition_conversion import (
     convert_project_definition_to_v2,
 )
@@ -293,3 +300,20 @@ def _validate_and_save_connections_imported_from_snowsql(
             path=["default_connection_name"],
             value=default_cli_connection_name,
         )
+
+
+@app.command(name="check-snowsql-env-vars", requires_connection=False)
+def check_snowsql_env_vars(**options):
+    """Check if there are any SnowSQL environment variables set."""
+
+    env_vars = os.environ.copy()
+    discovered, unused, summary = check_env_vars(env_vars)
+
+    results = []
+    if discovered:
+        results.append(CollectionResult(discovered))
+    if unused:
+        results.append(CollectionResult(unused))
+
+    results.append(MessageResult(summary))
+    return MultipleResults(results)

src/snowflake/cli/_plugins/helpers/snowsl_vars_reader.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+EnvVars = dict[str, str]
+
+DicoveredVars = list[dict[str, str]]
+UnusedVars = list[dict[str, str]]
+Summary = str
+
+CheckResult = tuple[DicoveredVars, UnusedVars, Summary]
+
+KNOWN_SNOWSQL_ENV_VARS = {
+    "SNOWSQL_ACCOUNT": {
+        "Found": "SNOWSQL_ACCOUNT",
+        "Suggested": "SNOWFLAKE_ACCOUNT",
+        "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
+    },
+    "SNOWSQL_PWD": {
+        "Found": "SNOWSQL_PASSWORD",
+        "Suggested": "SNOWFLAKE_PASSWORD",
+        "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
+    },
+    "SNOWSQL_USER": {
+        "Found": "SNOWSQL_USER",
+        "Suggested": "SNOWFLAKE_USER",
+        "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
+    },
+    "SNOWSQL_REGION": {
+        "Found": "SNOWSQL_REGION",
+        "Suggested": "SNOWFLAKE_REGION",
+        "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
+    },
+    "SNOWSQL_ROLE": {
+        "Found": "SNOWSQL_ROLE",
+        "Suggested": "SNOWFLAKE_ROLE",
+        "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
+    },
+    "SNOWSQL_WAREHOUSE": {
+        "Found": "SNOWSQL_WAREHOUSE",
+        "Suggested": "SNOWFLAKE_WAREHOUSE",
+        "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
+    },
+    "SNOWSQL_DATABASE": {
+        "Found": "SNOWSQL_DATABASE",
+        "Suggested": "SNOWFLAKE_DATABASE",
+        "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
+    },
+    "SNOWSQL_SCHEMA": {
+        "found": "SNOWSQL_SCHEMA",
+        "Suggested": "SNOWFLAKE_SCHEMA",
+        "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
+    },
+    "SNOWSQL_HOST": {
+        "Found": "SNOWSQL_HOST",
+        "Suggested": "SNOWFLAKE_HOST",
+        "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
+    },
+    "SNOWSQL_PORT": {
+        "Found": "SNOWSQL_PORT",
+        "Suggested": "SNOWFLAKE_PORT",
+        "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
+    },
+    "SNOWSQL_PROTOCOL": {
+        "Found": "SNOWSQL_PROTOCOL",
+        "Suggested": "SNOWFLAKE_PROTOCOL",
+        "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections#use-environment-variables-for-snowflake-credentials",
+    },
+    "SNOWSQL_PROXY_HOST": {
+        "Found": "SNOWSQL_PROXY_HOST",
+        "Suggested": "PROXY_HOST",
+        "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-cli#use-a-proxy-server",
+    },
+    "SNOWSQL_PROXY_PORT": {
+        "Found": "SNOWSQL_PROXY_HOST",
+        "Suggested": "PROXY_HOST",
+        "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-cli#use-a-proxy-server",
+    },
+    "SNOWSQL_PROXY_USER": {
+        "Found": "SNOWSQL_PROXY_PORT",
+        "Suggested": "PROXY_PORT",
+        "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-cli#use-a-proxy-server",
+    },
+    "SNOWSQL_PROXY_PWD": {
+        "Found": "SNOWSQL_PROXY_PWD",
+        "Suggested": "PROXY_PWD",
+        "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-cli#use-a-proxy-server",
+    },
+    "SNOWSQL_PRIVATE_KEY_PASSPHRASE": {
+        "Found": "SNOWSQL_PRIVATE_KEY_PASSPHRASE",
+        "Suggested": "PRIVATE_KEY_PASSPHRASE",
+        "Additional info": "https://docs.snowflake.com/en/developer-guide/snowflake-cli/connecting/configure-connections",
+    },
+}
+
+
+def check_env_vars(variables: EnvVars) -> CheckResult:
+    """Checks passed dict objects for possible SnowSQL variables.
+
+    Returns tuple of
+    - sequence of variables that can be adjusted
+    - sequence of variables that have no corresponding variables
+    - a summary messages
+    """
+    discovered: DicoveredVars = []
+    unused: UnusedVars = []
+
+    prefix_matched = (e for e in variables if e.lower().startswith("snowsql"))
+
+    for var in prefix_matched:
+        if suggestion := KNOWN_SNOWSQL_ENV_VARS.get(var, None):
+            discovered.append(suggestion)
+        else:
+            unused.append(
+                {
+                    "Found": var,
+                    "Suggested": "n/a",
+                    "Additional info": "Unused variable",
+                }
+            )
+
+    discovered_cnt = len(discovered)
+    unused_cnt = len(unused)
+
+    summary: Summary = (
+        f"Found {discovered_cnt + unused_cnt} SnowSQL environment variables,"
+        f" {discovered_cnt} with replacements, {unused_cnt} unused."
+    )
+
+    return discovered, unused, summary

tests/helpers/test_snowsql_env_vars.py

@@ -0,0 +1,42 @@
+import os
+from unittest import mock
+
+COMMAND = "check-snowsql-env-vars"
+
+VALID_VARS = {
+    "SNOWSQL_ACCOUNT": "test",
+    "SNOWSQL_WAREHOUSE": "whwhw",
+}
+INVALID_VARS = {
+    "SNOWSQL_ABC": "test",
+    "SNOWSQL_CBD": "whwhw",
+}
+
+
+@mock.patch.dict(os.environ, VALID_VARS, clear=True)
+def test_replecaple_env_found(runner):
+    result = runner.invoke(("helpers", COMMAND))
+    assert result.exit_code == 0
+    assert (
+        "Found 2 SnowSQL environment variables, 2 with replacements, 0 unused."
+        in result.output
+    )
+
+
+@mock.patch.dict(os.environ, INVALID_VARS, clear=True)
+def test_non_replacable_env_found(runner):
+    result = runner.invoke(("helpers", COMMAND))
+    assert result.exit_code == 0
+    assert (
+        "Found 2 SnowSQL environment variables, 0 with replacements, 2 unused."
+        in result.output
+    )
+
+
+def test_no_env_found(runner):
+    result = runner.invoke(("helpers", COMMAND))
+    assert result.exit_code == 0
+    assert (
+        "Found 0 SnowSQL environment variables, 0 with replacements, 0 unused."
+        in result.output
+    )