diff --git a/changelog.md b/changelog.md index dd3aa75f..21c296a8 100644 --- a/changelog.md +++ b/changelog.md @@ -4,6 +4,7 @@ Upcoming (TBD) Features --------- * Add `--ssh-options` CLI argument to pass extra options with `--ssh-jump`. +* Make `ssh_jump` a DSN query parameter. Bug Fixes diff --git a/mycli/cli_runner.py b/mycli/cli_runner.py index 3c61390f..f03b71c9 100644 --- a/mycli/cli_runner.py +++ b/mycli/cli_runner.py @@ -254,6 +254,8 @@ def run_from_cli_args(cli_args: 'CliArgs', client_factory: ClientFactory) -> Non cli_args.keepalive_ticks = int(params[0]) if params := dsn_params.get('character_set'): cli_args.character_set = cli_args.character_set or params[0] + if params := dsn_params.get('ssh_jump'): + cli_args.ssh_jump = cli_args.ssh_jump or params[0] keepalive_ticks = cli_args.keepalive_ticks if cli_args.keepalive_ticks is not None else mycli.default_keepalive_ticks ssl_mode = cli_args.ssl_mode or mycli.ssl_mode diff --git a/mycli/client_connection.py b/mycli/client_connection.py index 9558b943..300b5e15 100644 --- a/mycli/client_connection.py +++ b/mycli/client_connection.py @@ -198,6 +198,7 @@ def connect( socket=remote_socket, database=database, character_set=character_set, + ssh_jump=ssh_jump, ) connection_info: dict[str, Any] = { diff --git a/mycli/packages/special/utils.py b/mycli/packages/special/utils.py index 9b78c399..fe422788 100644 --- a/mycli/packages/special/utils.py +++ b/mycli/packages/special/utils.py @@ -155,6 +155,7 @@ def format_connection_dsn( database: str | None, socket: str | None, character_set: str | None, + ssh_jump: str | None = None, ) -> str: user = urlquote(user or '') host = host or 'localhost' @@ -168,6 +169,8 @@ def format_connection_dsn( port_part = '' if character_set and character_set != 'utf8mb4': query_part['character_set'] = character_set + if ssh_jump: + query_part['ssh_jump'] = ssh_jump dsn = f'mysql://{user}@{host}{port_part}{db}' if query_part: dsn += '?' + urlencode(query_part) diff --git a/test/pytests/test_cli_runner.py b/test/pytests/test_cli_runner.py index 1af17c6f..efba80a2 100644 --- a/test/pytests/test_cli_runner.py +++ b/test/pytests/test_cli_runner.py @@ -217,6 +217,22 @@ def test_run_from_cli_args_expands_whole_dsn_alias_env_vars_when_enabled( assert client.connect_calls[-1]['keepalive_ticks'] == 9 +def test_run_from_cli_args_expands_dsn_alias_ssh_jump_env_var_when_enabled( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cli_args = make_cli_args() + cli_args.dsn = 'prod' + monkeypatch.setenv('MYCLI_TEST_DSN_SSH_JUMP', 'env-bastion') + config = default_config() + config['main'] = {**config['main'], 'expand_dsn_alias_env_vars': 'true'} + config['alias_dsn'] = {'prod': 'mysql://user@host/db?ssh_jump=${MYCLI_TEST_DSN_SSH_JUMP}'} + client = DummyMyCli(config=config) + + run_with_client(monkeypatch, cli_args, client) + + assert client.connect_calls[-1]['ssh_jump'] == 'env-bastion' + + def test_run_from_cli_args_does_not_expand_partial_values_or_query_keys( monkeypatch: pytest.MonkeyPatch, ) -> None: @@ -354,6 +370,29 @@ def test_run_from_cli_args_accepts_mysql_plus_dsn_scheme(monkeypatch: pytest.Mon assert client.connect_calls[-1]['database'] == 'db' +def test_run_from_cli_args_maps_dsn_ssh_jump_parameter(monkeypatch: pytest.MonkeyPatch) -> None: + cli_args = make_cli_args() + cli_args.dsn = 'mysql://user@host/db?ssh_jump=bastion' + client = DummyMyCli() + + run_with_client(monkeypatch, cli_args, client) + + assert client.connect_calls[-1]['ssh_jump'] == 'bastion' + + +def test_run_from_cli_args_prefers_cli_ssh_jump_over_dsn_parameter( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cli_args = make_cli_args() + cli_args.dsn = 'mysql://user@host/db?ssh_jump=dsn-bastion' + cli_args.ssh_jump = 'cli-bastion' + client = DummyMyCli() + + run_with_client(monkeypatch, cli_args, client) + + assert client.connect_calls[-1]['ssh_jump'] == 'cli-bastion' + + def test_run_from_cli_args_maps_dsn_ssl_parameters(monkeypatch: pytest.MonkeyPatch) -> None: cli_args = make_cli_args() cli_args.dsn = ( diff --git a/test/pytests/test_client_connection.py b/test/pytests/test_client_connection.py index e8543506..5bdfd5cb 100644 --- a/test/pytests/test_client_connection.py +++ b/test/pytests/test_client_connection.py @@ -384,7 +384,7 @@ def close(self) -> None: assert FakeSQLExecute.calls[-1]['host'] is None assert FakeSQLExecute.calls[-1]['port'] is None assert FakeSQLExecute.calls[-1]['socket'] == '/tmp/mycli-ssh.sock' - assert FakeSQLExecute.calls[-1]['display_dsn'] == 'mysql://alice@localhost?socket=%2Fvar%2Frun%2Fmysqld%2Fmysqld.sock' + assert FakeSQLExecute.calls[-1]['display_dsn'] == 'mysql://alice@localhost?socket=%2Fvar%2Frun%2Fmysqld%2Fmysqld.sock&ssh_jump=bastion' def test_connect_uses_ssh_jump_with_local_port(monkeypatch: pytest.MonkeyPatch) -> None: @@ -422,7 +422,7 @@ def close(self) -> None: assert FakeSQLExecute.calls[-1]['host'] == '127.0.0.1' assert FakeSQLExecute.calls[-1]['port'] == 4406 assert FakeSQLExecute.calls[-1]['socket'] is None - assert FakeSQLExecute.calls[-1]['display_dsn'] == 'mysql://alice@db.internal:3307' + assert FakeSQLExecute.calls[-1]['display_dsn'] == 'mysql://alice@db.internal:3307?ssh_jump=bastion' def test_connect_passes_config_and_cli_ssh_options(monkeypatch: pytest.MonkeyPatch) -> None: diff --git a/test/pytests/test_special_utils.py b/test/pytests/test_special_utils.py index 6238d1da..7edf1a03 100644 --- a/test/pytests/test_special_utils.py +++ b/test/pytests/test_special_utils.py @@ -13,6 +13,7 @@ from mycli.packages.special.utils import ( CACHED_SSL_VERSION, compute_current_dsn, + format_connection_dsn, format_uptime, get_local_timezone, get_server_timezone, @@ -311,6 +312,21 @@ def test_compute_current_dsn_prefers_mycli_display_dsn() -> None: assert compute_current_dsn(cursor) == 'mysql://alice@db.example.com:3307/prod' +def test_format_connection_dsn_includes_ssh_jump() -> None: + assert ( + format_connection_dsn( + user='alice', + host='db.example.com', + port=3307, + database='prod', + socket=None, + character_set='utf8mb4', + ssh_jump='bastion', + ) + == 'mysql://alice@db.example.com:3307/prod?ssh_jump=bastion' + ) + + def test_compute_current_dsn_for_socket_connection() -> None: connection = SimpleNamespace( user='alice',