From 6ad5ca86611d478d2f9a4afe36101d101d711eb7 Mon Sep 17 00:00:00 2001 From: Tomasz Urbaszek Date: Mon, 14 Oct 2024 15:38:22 +0200 Subject: [PATCH] Add connection generate-jwt command (#1712) --- RELEASE-NOTES.md | 1 + .../cli/_plugins/connection/commands.py | 49 ++++++++++++++++ tests/__snapshots__/test_help_messages.ambr | 51 ++++++++++++++--- tests/test_connection.py | 57 +++++++++++++++++++ 4 files changed, 150 insertions(+), 8 deletions(-) diff --git a/RELEASE-NOTES.md b/RELEASE-NOTES.md index 4d0ecaddaa..72de5abb5c 100644 --- a/RELEASE-NOTES.md +++ b/RELEASE-NOTES.md @@ -19,6 +19,7 @@ ## Deprecations ## New additions +* Added `snow connection generate-jwt` command to generate JWT token for Snowflake connection. ## Fixes and improvements * Fixed a bug that would cause the `deploy_root`, `bundle_root`, and `generated_root` directories to be created in the current working directory instead of the project root when invoking commands with the `--project` flag from a different directory. diff --git a/src/snowflake/cli/_plugins/connection/commands.py b/src/snowflake/cli/_plugins/connection/commands.py index 879b75e93d..543f70f285 100644 --- a/src/snowflake/cli/_plugins/connection/commands.py +++ b/src/snowflake/cli/_plugins/connection/commands.py @@ -16,11 +16,13 @@ import logging import os.path +from pathlib import 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 import connector from snowflake.cli._plugins.connection.util import ( strip_and_check_if_exists, strip_if_value_present, @@ -351,3 +353,50 @@ def set_default( 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}") + + +@app.command(requires_connection=False) +def generate_jwt( + account: str = typer.Option( + None, + "--account", + "-a", + "--accountname", + help="Account name to use when authenticating with Snowflake.", + show_default=False, + ), + user: str = typer.Option( + None, + "--user", + "-u", + "--username", + show_default=False, + help="Username to connect to Snowflake.", + ), + private_key_file: Path = typer.Option( + None, + "--private-key", + "--private-key-path", + "-k", + help="Path to file containing private key", + dir_okay=False, + exists=True, + ), + **options, +) -> CommandResult: + """Generate a JWT token, which will be printed out and displayed..""" + passphrase = os.getenv("PRIVATE_KEY_PASSPHRASE", None) + if not passphrase: + passphrase = typer.prompt( + "Enter private key file password (Press enter if none)", + hide_input=True, + type=str, + default="", + ) + try: + token = connector.auth.get_token_from_private_key( + user, account, private_key_file, passphrase + ) + return MessageResult(token) + except ValueError as err: + raise ClickException(str(err)) diff --git a/tests/__snapshots__/test_help_messages.ambr b/tests/__snapshots__/test_help_messages.ambr index 597d76d6a4..1b9f519075 100644 --- a/tests/__snapshots__/test_help_messages.ambr +++ b/tests/__snapshots__/test_help_messages.ambr @@ -1504,6 +1504,37 @@ +------------------------------------------------------------------------------+ + ''' +# --- +# name: test_help_messages[connection.generate-jwt] + ''' + + Usage: default connection generate-jwt [OPTIONS] + + Generate a JWT token, which will be printed out and displayed.. + + +- Options --------------------------------------------------------------------+ + | --account,--accountname -a TEXT Account name to use when | + | authenticating with Snowflake. | + | --user,--username -u TEXT Username to connect to | + | Snowflake. | + | --private-key,--private-key-p… -k FILE Path to file containing | + | private key | + | [default: None] | + | --help -h Show this message and exit. | + +------------------------------------------------------------------------------+ + +- Global configuration -------------------------------------------------------+ + | --format [TABLE|JSON] Specifies the output format. | + | [default: TABLE] | + | --verbose -v Displays log entries for log levels info | + | and higher. | + | --debug Displays log entries for log levels debug | + | and higher; debug logs contains additional | + | information. | + | --silent Turns off intermediate output to console. | + +------------------------------------------------------------------------------+ + + ''' # --- # name: test_help_messages[connection.list] @@ -1654,10 +1685,12 @@ | --help -h Show this message and exit. | +------------------------------------------------------------------------------+ +- Commands -------------------------------------------------------------------+ - | add Adds a connection to configuration file. | - | list Lists configured connections. | - | set-default Changes default connection to provided value. | - | test Tests the connection to Snowflake. | + | add Adds a connection to configuration file. | + | generate-jwt Generate a JWT token, which will be printed out and | + | displayed.. | + | list Lists configured connections. | + | set-default Changes default connection to provided value. | + | test Tests the connection to Snowflake. | +------------------------------------------------------------------------------+ @@ -9423,10 +9456,12 @@ | --help -h Show this message and exit. | +------------------------------------------------------------------------------+ +- Commands -------------------------------------------------------------------+ - | add Adds a connection to configuration file. | - | list Lists configured connections. | - | set-default Changes default connection to provided value. | - | test Tests the connection to Snowflake. | + | add Adds a connection to configuration file. | + | generate-jwt Generate a JWT token, which will be printed out and | + | displayed.. | + | list Lists configured connections. | + | set-default Changes default connection to provided value. | + | test Tests the connection to Snowflake. | +------------------------------------------------------------------------------+ diff --git a/tests/test_connection.py b/tests/test_connection.py index 657fd9b454..127c20271e 100644 --- a/tests/test_connection.py +++ b/tests/test_connection.py @@ -1135,3 +1135,60 @@ def test_new_connection_is_added_to_connections_toml( port = "8080" """ ) + + +@mock.patch( + "snowflake.cli._plugins.connection.commands.connector.auth.get_token_from_private_key" +) +def test_generate_jwt(mocked_get_token, runner, named_temporary_file): + mocked_get_token.return_value = "funny token" + + with named_temporary_file() as f: + f.write_text("secret from file") + result = runner.invoke( + [ + "connection", + "generate-jwt", + "--user", + "FooBar", + "--account", + "account1", + "--private-key-path", + f, + ], + input="123", + ) + + assert result.exit_code == 0, result.output + assert ( + result.output + == "Enter private key file password (Press enter if none) []: \nfunny token\n" + ) + mocked_get_token.assert_called_once_with("FooBar", "account1", f, "123") + + +@mock.patch.dict(os.environ, {"PRIVATE_KEY_PASSPHRASE": "123"}) +@mock.patch( + "snowflake.cli._plugins.connection.commands.connector.auth.get_token_from_private_key" +) +def test_generate_jwt_with_pass_phrase(mocked_get_token, runner, named_temporary_file): + mocked_get_token.return_value = "funny token" + + with named_temporary_file() as f: + f.write_text("secret from file") + result = runner.invoke( + [ + "connection", + "generate-jwt", + "--user", + "FooBar", + "--account", + "account1", + "--private-key-path", + f, + ] + ) + + assert result.exit_code == 0, result.output + assert result.output == "funny token\n" + mocked_get_token.assert_called_once_with("FooBar", "account1", f, "123")