Skip to content

Commit

Permalink
Add connection generate-jwt command (#1712)
Browse files Browse the repository at this point in the history
  • Loading branch information
@sfc-gh-turbaszek
sfc-gh-turbaszek authored Oct 14, 2024
1 parent 9232ed6 commit 6ad5ca8
Show file tree
Hide file tree
Showing 4 changed files with 150 additions and 8 deletions.
1 change: 1 addition & 0 deletions RELEASE-NOTES.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
49 changes: 49 additions & 0 deletions src/snowflake/cli/_plugins/connection/commands.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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))
51 changes: 43 additions & 8 deletions tests/__snapshots__/test_help_messages.ambr
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down Expand Up @@ -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. |
+------------------------------------------------------------------------------+


Expand Down Expand Up @@ -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. |
+------------------------------------------------------------------------------+


Expand Down
57 changes: 57 additions & 0 deletions tests/test_connection.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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")

0 comments on commit 6ad5ca8

Please sign in to comment.