Make ssh_jump a DSN query parameter

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

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

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] = {

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)

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 = (

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:

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',