Add SSH tunnel support (#1301)

* Add initial sshtunnel support

* Force CI to rerun.

* Fix unit test for 3.6.

* Black.

Co-authored-by: Irina Truong <i.chernyavska@gmail.com>

2 files changed, 133 insertions, 2 deletions

diff --git a/tests/features/steps/basic_commands.py b/tests/features/steps/basic_commands.py
index 7ca20f06..a7c99eea 100644
--- a/tests/features/steps/basic_commands.py
+++ b/tests/features/steps/basic_commands.py
@@ -97,9 +97,9 @@ def step_see_error_message(context):
 @when("we send source command")
 def step_send_source_command(context):
     context.tmpfile_sql_help = tempfile.NamedTemporaryFile(prefix="pgcli_")
-    context.tmpfile_sql_help.write(br"\?")
+    context.tmpfile_sql_help.write(rb"\?")
     context.tmpfile_sql_help.flush()
-    context.cli.sendline(fr"\i {context.tmpfile_sql_help.name}")
+    context.cli.sendline(rf"\i {context.tmpfile_sql_help.name}")
     wrappers.expect_exact(context, context.conf["pager_boundary"] + "\r\n", timeout=5)
 
 
diff --git a/tests/test_ssh_tunnel.py b/tests/test_ssh_tunnel.py
new file mode 100644
index 00000000..9163c08c
--- /dev/null
+++ b/tests/test_ssh_tunnel.py
@@ -0,0 +1,131 @@
+from unittest.mock import patch, MagicMock, ANY
+
+import pytest
+from click.testing import CliRunner
+from sshtunnel import SSHTunnelForwarder
+
+from pgcli.main import cli, PGCli
+from pgcli.pgexecute import PGExecute
+
+
+@pytest.fixture
+def mock_ssh_tunnel_forwarder() -> MagicMock:
+    mock_ssh_tunnel_forwarder = MagicMock(
+        SSHTunnelForwarder, local_bind_ports=[1111], autospec=True
+    )
+    with patch(
+        "pgcli.main.sshtunnel.SSHTunnelForwarder",
+        return_value=mock_ssh_tunnel_forwarder,
+    ) as mock:
+        yield mock
+
+
+@pytest.fixture
+def mock_pgexecute() -> MagicMock:
+    with patch.object(PGExecute, "__init__", return_value=None) as mock_pgexecute:
+        yield mock_pgexecute
+
+
+def test_ssh_tunnel(
+    mock_ssh_tunnel_forwarder: MagicMock, mock_pgexecute: MagicMock
+) -> None:
+    # Test with just a host
+    tunnel_url = "some.host"
+    db_params = {
+        "database": "dbname",
+        "host": "db.host",
+        "user": "db_user",
+        "passwd": "db_passwd",
+    }
+    expected_tunnel_params = {
+        "local_bind_address": ("127.0.0.1",),
+        "remote_bind_address": (db_params["host"], 5432),
+        "ssh_address_or_host": (tunnel_url, 22),
+        "logger": ANY,
+    }
+
+    pgcli = PGCli(ssh_tunnel_url=tunnel_url)
+    pgcli.connect(**db_params)
+
+    mock_ssh_tunnel_forwarder.assert_called_once_with(**expected_tunnel_params)
+    mock_ssh_tunnel_forwarder.return_value.start.assert_called_once()
+    mock_pgexecute.assert_called_once()
+
+    call_args, call_kwargs = mock_pgexecute.call_args
+    assert call_args == (
+        db_params["database"],
+        db_params["user"],
+        db_params["passwd"],
+        "127.0.0.1",
+        pgcli.ssh_tunnel.local_bind_ports[0],
+        "",
+    )
+    mock_ssh_tunnel_forwarder.reset_mock()
+    mock_pgexecute.reset_mock()
+
+    # Test with a full url and with a specific db port
+    tunnel_user = "tunnel_user"
+    tunnel_passwd = "tunnel_pass"
+    tunnel_host = "some.other.host"
+    tunnel_port = 1022
+    tunnel_url = f"ssh://{tunnel_user}:{tunnel_passwd}@{tunnel_host}:{tunnel_port}"
+    db_params["port"] = 1234
+
+    expected_tunnel_params["remote_bind_address"] = (
+        db_params["host"],
+        db_params["port"],
+    )
+    expected_tunnel_params["ssh_address_or_host"] = (tunnel_host, tunnel_port)
+    expected_tunnel_params["ssh_username"] = tunnel_user
+    expected_tunnel_params["ssh_password"] = tunnel_passwd
+
+    pgcli = PGCli(ssh_tunnel_url=tunnel_url)
+    pgcli.connect(**db_params)
+
+    mock_ssh_tunnel_forwarder.assert_called_once_with(**expected_tunnel_params)
+    mock_ssh_tunnel_forwarder.return_value.start.assert_called_once()
+    mock_pgexecute.assert_called_once()
+
+    call_args, call_kwargs = mock_pgexecute.call_args
+    assert call_args == (
+        db_params["database"],
+        db_params["user"],
+        db_params["passwd"],
+        "127.0.0.1",
+        pgcli.ssh_tunnel.local_bind_ports[0],
+        "",
+    )
+    mock_ssh_tunnel_forwarder.reset_mock()
+    mock_pgexecute.reset_mock()
+
+    # Test with DSN
+    dsn = (
+        f"user={db_params['user']} password={db_params['passwd']} "
+        f"host={db_params['host']} port={db_params['port']}"
+    )
+
+    pgcli = PGCli(ssh_tunnel_url=tunnel_url)
+    pgcli.connect(dsn=dsn)
+
+    expected_dsn = (
+        f"user={db_params['user']} password={db_params['passwd']} "
+        f"host=127.0.0.1 port={pgcli.ssh_tunnel.local_bind_ports[0]}"
+    )
+
+    mock_ssh_tunnel_forwarder.assert_called_once_with(**expected_tunnel_params)
+    mock_pgexecute.assert_called_once()
+
+    call_args, call_kwargs = mock_pgexecute.call_args
+    assert expected_dsn in call_args
+
+
+def test_cli_with_tunnel() -> None:
+    runner = CliRunner()
+    tunnel_url = "mytunnel"
+    with patch.object(
+        PGCli, "__init__", autospec=True, return_value=None
+    ) as mock_pgcli:
+        runner.invoke(cli, ["--ssh-tunnel", tunnel_url])
+        mock_pgcli.assert_called_once()
+        call_args, call_kwargs = mock_pgcli.call_args
+        assert call_kwargs["ssh_tunnel_url"] == tunnel_url