Reuse ssh connection

changelog.md

@@ -13,6 +13,7 @@ Features:
 * Add `-g` shortcut to option `--login-path`.
 * Alt-Enter dispatches the command in multi-line mode.
 * Allow to pass a file or FIFO path with --password-file when password is not specified or is failing (as suggested in this best-practice https://www.netmeister.org/blog/passing-passwords.html)
+* Reuse the same SSH connection for both main thread and completion thread.
 
 Internal:
 ---------
@@ -846,6 +847,7 @@ Bug Fixes:
 [Georgy Frolov]: https://github.com/pasenor
 [Zach DeCook]: https://zachdecook.com
 [laixintao]: https://github.com/laixintao
+[Nathan Huang]: https://github.com/hxueh
 [mtorromeo]: https://github.com/mtorromeo
 [mwcm]: https://github.com/mwcm
 [xeron]: https://github.com/xeron

mycli/completion_refresher.py

@@ -51,8 +51,7 @@ def _bg_refresh(self, sqlexecute, callbacks, completer_options):
  e = sqlexecute
  executor = SQLExecute(e.dbname, e.user, e.password, e.host, e.port,
  e.socket, e.charset, e.local_infile, e.ssl,
- e.ssh_user, e.ssh_host, e.ssh_port,
- e.ssh_password, e.ssh_key_filename)
+ ssh_client=e.ssh_client)
 
  # If callbacks is a single function then push it into a list.
  if callable(callbacks):

mycli/main.py

@@ -37,10 +37,11 @@
 from prompt_toolkit.history import FileHistory
 from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
 
+from mycli.packages.ssh_client import create_ssh_client
 from .packages.special.main import NO_QUERY
 from .packages.prompt_utils import confirm, confirm_destructive_query
 from .packages.tabular_output import sql_format
-from .packages import special
+from .packages import special, ssh_client
 from .packages.special.favoritequeries import FavoriteQueries
 from .sqlcompleter import SQLCompleter
 from .clitoolbar import create_toolbar_tokens_func
@@ -74,11 +75,6 @@
  # Python < 3.7
  import importlib_resources as resources
 
-try:
- import paramiko
-except ImportError:
- from mycli.packages.paramiko_stub import paramiko
-
 # Query tuples are used for maintaining history
 Query = namedtuple('Query', ['query', 'successful', 'mutating'])
 
@@ -211,6 +207,8 @@ def __init__(self, sqlexecute=None, prompt=None,
 
  self.prompt_app = None
 
+ self.ssh_client = None
+
  def register_special_commands(self):
  special.register_special_command(self.change_db, 'use',
  '\\u', 'Change to a new database.', aliases=('\\u',))
@@ -387,9 +385,8 @@ def merge_ssl_with_cnf(self, ssl, cnf):
  return merged
 
  def connect(self, database='', user='', passwd='', host='', port='',
- socket='', charset='', local_infile='', ssl='',
- ssh_user='', ssh_host='', ssh_port='',
- ssh_password='', ssh_key_filename='', init_command='', password_file=''):
+ socket='', charset='', local_infile='', ssl=None, init_command='',
+ password_file=''):
 
  cnf = {'database': None,
  'user': None,
@@ -418,7 +415,7 @@ def connect(self, database='', user='', passwd='', host='', port='',
  ssl = ssl or {}
 
  port = port and int(port)
- if not port:
+ if not port and not self.ssh_client:
  port = 3306
  if not host or host == 'localhost':
  socket = (
@@ -455,8 +452,7 @@ def _connect():
  try:
  self.sqlexecute = SQLExecute(
  database, user, passwd, host, port, socket, charset,
- local_infile, ssl, ssh_user, ssh_host, ssh_port,
- ssh_password, ssh_key_filename, init_command
+ local_infile, ssl, init_command, ssh_client=self.ssh_client
  )
  except OperationalError as e:
  if e.args[0] == ERROR_CODE_ACCESS_DENIED:
@@ -467,8 +463,8 @@ def _connect():
  show_default=False, type=str, err=True)
  self.sqlexecute = SQLExecute(
  database, user, new_passwd, host, port, socket,
- charset, local_infile, ssl, ssh_user, ssh_host,
- ssh_port, ssh_password, ssh_key_filename, init_command
+ charset, local_infile, ssl, init_command,
+ ssh_client=self.ssh_client
  )
  else:
  raise e
@@ -1179,16 +1175,22 @@ def cli(database, user, host, port, socket, password, dbname,
  else:
  click.secho(alias)
  sys.exit(0)
+
  if list_ssh_config:
- ssh_config = read_ssh_config(ssh_config_path)
- for host in ssh_config.get_hostnames():
+ try:
+ hosts = ssh_client.get_config_hosts(ssh_config_path)
+ except ssh_client.SSHException as e:
+ click.secho(str(e), err=True, fg='red')
+ sys.exit(1)
+
+ for host, hostname in hosts.items():
  if verbose:
- host_config = ssh_config.lookup(host)
  click.secho("{} : {}".format(
- host, host_config.get('hostname')))
+ host, hostname))
  else:
  click.secho(host)
  sys.exit(0)
+
  # Choose which ever one has a valid value.
  database = dbname or database
 
@@ -1240,9 +1242,14 @@ def cli(database, user, host, port, socket, password, dbname,
  port = uri.port
 
  if ssh_config_host:
- ssh_config = read_ssh_config(
- ssh_config_path
- ).lookup(ssh_config_host)
+ try:
+ ssh_config = ssh_client.read_config_file(
+ ssh_config_path
+ ).lookup(ssh_config_host)
+ except ssh_client.SSHException as e:
+ click.secho(str(e), err=True, fg='red')
+ sys.exit(1)
+
  ssh_host = ssh_host if ssh_host else ssh_config.get('hostname')
  ssh_user = ssh_user if ssh_user else ssh_config.get('user')
  if ssh_config.get('port') and ssh_port == 22:
@@ -1251,7 +1258,10 @@ def cli(database, user, host, port, socket, password, dbname,
  ssh_key_filename = ssh_key_filename if ssh_key_filename else ssh_config.get(
  'identityfile', [None])[0]
 
- ssh_key_filename = ssh_key_filename and os.path.expanduser(ssh_key_filename)
+ if ssh_host:
+ mycli.ssh_client = create_ssh_client(
+ ssh_host, ssh_port, ssh_user, ssh_password, ssh_key_filename
+ )
 
  mycli.connect(
  database=database,
@@ -1262,14 +1272,9 @@ def cli(database, user, host, port, socket, password, dbname,
  socket=socket,
  local_infile=local_infile,
  ssl=ssl,
- ssh_user=ssh_user,
- ssh_host=ssh_host,
- ssh_port=ssh_port,
- ssh_password=ssh_password,
- ssh_key_filename=ssh_key_filename,
  init_command=init_command,
  charset=charset,
- password_file=password_file
+ password_file=password_file,
  )
 
  mycli.logger.debug('Launch Params: \n'
@@ -1403,6 +1408,7 @@ def read_ssh_config(ssh_config_path):
  except FileNotFoundError as e:
  click.secho(str(e), err=True, fg='red')
  sys.exit(1)
+
  # Paramiko prior to version 2.7 raises Exception on parse errors.
  # In 2.7 it has become paramiko.ssh_exception.SSHException,
  # but let's catch everything for compatibility

mycli/packages/ssh_client/__init__.py

@@ -0,0 +1 @@
+from .client import get_config_hosts, create_ssh_client, SSHException, read_config_file

mycli/packages/ssh_client/client.py

@@ -0,0 +1,57 @@
+"""A very thin wrapper around paramiko, mostly to keep all SSH-related
+functionality in one place."""
+from io import open
+import logging
+
+_logger = logging.getLogger(__name__)
+
+try:
+ import paramiko
+except ImportError:
+ from mycli.packages.paramiko_stub import paramiko
+
+
+class SSHException(Exception):
+ pass
+
+
+def get_config_hosts(config_path):
+ config = read_config_file(config_path)
+ return {
+ host: config.lookup(host).get("hostname") for host in config.get_hostnames()
+ }
+
+
+def create_ssh_client(ssh_host, ssh_port, ssh_user, ssh_password=None, ssh_key_filename=None) -> paramiko.SSHClient:
+ client = paramiko.SSHClient()
+ client.load_system_host_keys()
+ client.set_missing_host_key_policy(paramiko.WarningPolicy())
+ _logger.debug(
+ f'Connecting to ssh server with \n'
+ ' host = {ssh_host}\n'
+ ' port = {ssh_port}\n'
+ ' user = {ssh_user}\n'
+ ' password = {ssh_password}\n'
+ ' key_filename = {ssh_key_filename}\n'
+ )
+ client.connect(
+ ssh_host, ssh_port, ssh_user, password=ssh_password, key_filename=ssh_key_filename
+ )
+ return client
+
+
+def read_config_file(config_path) -> paramiko.SSHConfig:
+ ssh_config = paramiko.config.SSHConfig()
+ try:
+ with open(config_path) as f:
+ ssh_config.parse(f)
+ except FileNotFoundError as e:
+ raise SSHException(str(e))
+ # Paramiko prior to version 2.7 raises Exception on parse errors.
+ # In 2.7 it has become paramiko.ssh_exception.SSHException,
+ # but let's catch everything for compatibility
+ except Exception as err:
+ raise SSHException(
+ f"Could not parse SSH configuration file {config_path}:\n{err} ",
+ )
+ return ssh_config

mycli/sqlexecute.py

@@ -8,10 +8,7 @@
 from pymysql.converters import (convert_datetime,
  convert_timedelta, convert_date, conversions,
  decoders)
-try:
- import paramiko
-except ImportError:
- from mycli.packages.paramiko_stub import paramiko
+
 
 _logger = logging.getLogger(__name__)
 
@@ -97,8 +94,8 @@ class SQLExecute(object):
  order by table_name,ordinal_position'''
 
  def __init__(self, database, user, password, host, port, socket, charset,
- local_infile, ssl, ssh_user, ssh_host, ssh_port, ssh_password,
- ssh_key_filename, init_command=None):
+ local_infile, ssl, init_command=None,
+ ssh_client=None):
  self.dbname = database
  self.user = user
  self.password = password
@@ -110,18 +107,14 @@ def __init__(self, database, user, password, host, port, socket, charset,
  self.ssl = ssl
  self.server_info = None
  self.connection_id = None
- self.ssh_user = ssh_user
- self.ssh_host = ssh_host
- self.ssh_port = ssh_port
- self.ssh_password = ssh_password
- self.ssh_key_filename = ssh_key_filename
  self.init_command = init_command
+ self.ssh_client = ssh_client
+
  self.connect()
 
  def connect(self, database=None, user=None, password=None, host=None,
- port=None, socket=None, charset=None, local_infile=None,
- ssl=None, ssh_host=None, ssh_port=None, ssh_user=None,
- ssh_password=None, ssh_key_filename=None, init_command=None):
+ port=None, socket=None, charset=None, local_infile=None, ssl=None,
+ ssh_client=None, init_command=None):
  db = (database or self.dbname)
  user = (user or self.user)
  password = (password or self.password)
@@ -131,11 +124,6 @@ def connect(self, database=None, user=None, password=None, host=None,
  charset = (charset or self.charset)
  local_infile = (local_infile or self.local_infile)
  ssl = (ssl or self.ssl)
- ssh_user = (ssh_user or self.ssh_user)
- ssh_host = (ssh_host or self.ssh_host)
- ssh_port = (ssh_port or self.ssh_port)
- ssh_password = (ssh_password or self.ssh_password)
- ssh_key_filename = (ssh_key_filename or self.ssh_key_filename)
  init_command = (init_command or self.init_command)
  _logger.debug(
  'Connection DB Params: \n'
@@ -147,15 +135,10 @@ def connect(self, database=None, user=None, password=None, host=None,
  '\tcharset: %r'
  '\tlocal_infile: %r'
  '\tssl: %r'
- '\tssh_user: %r'
- '\tssh_host: %r'
- '\tssh_port: %r'
- '\tssh_password: %r'
- '\tssh_key_filename: %r'
- '\tinit_command: %r',
+ '\tinit_command: %r'
+ '\tusing ssh: %r',
  db, user, host, port, socket, charset, local_infile, ssl,
- ssh_user, ssh_host, ssh_port, ssh_password, ssh_key_filename,
- init_command
+ init_command, bool(ssh_client)
  )
  conv = conversions.copy()
  conv.update({
@@ -167,9 +150,6 @@ def connect(self, database=None, user=None, password=None, host=None,
 
  defer_connect = False
 
- if ssh_host:
- defer_connect = True
-
  client_flag = pymysql.constants.CLIENT.INTERACTIVE
  if init_command and len(list(special.split_queries(init_command))) > 1:
  client_flag |= pymysql.constants.CLIENT.MULTI_STATEMENTS
@@ -179,18 +159,12 @@ def connect(self, database=None, user=None, password=None, host=None,
  unix_socket=socket, use_unicode=True, charset=charset,
  autocommit=True, client_flag=client_flag,
  local_infile=local_infile, conv=conv, ssl=ssl, program_name="mycli",
- defer_connect=defer_connect, init_command=init_command
+ defer_connect=self.ssh_client is not None,
+ init_command=init_command
  )
 
- if ssh_host:
- client = paramiko.SSHClient()
- client.load_system_host_keys()
- client.set_missing_host_key_policy(paramiko.WarningPolicy())
- client.connect(
- ssh_host, ssh_port, ssh_user, ssh_password,
- key_filename=ssh_key_filename
- )
- chan = client.get_transport().open_channel(
+ if self.ssh_client:
+ chan = self.ssh_client.get_transport().open_channel(
  'direct-tcpip',
  (host, port),
  ('0.0.0.0', 0),

test/conftest.py

@@ -1,4 +1,6 @@
 import pytest
+
+from mycli.packages.ssh_client import create_ssh_client
 from .utils import (HOST, USER, PASSWORD, PORT, CHARSET, create_db,
  db_connection, SSH_USER, SSH_HOST, SSH_PORT)
 import mycli.sqlexecute
@@ -21,9 +23,13 @@ def cursor(connection):
 
 @pytest.fixture
 def executor(connection):
+ if SSH_HOST:
+ ssh_client = create_ssh_client(SSH_HOST, SSH_PORT, SSH_USER)
+ else:
+ ssh_client = None
+
  return mycli.sqlexecute.SQLExecute(
  database='_test_db', user=USER,
  host=HOST, password=PASSWORD, port=PORT, socket=None, charset=CHARSET,
- local_infile=False, ssl=None, ssh_user=SSH_USER, ssh_host=SSH_HOST,
- ssh_port=SSH_PORT, ssh_password=None, ssh_key_filename=None
+ local_infile=False, ssl=None, ssh_client=ssh_client
  )