diff --git a/connections.py b/connections.py index e20208a..1e364b9 100644 --- a/connections.py +++ b/connections.py @@ -124,3 +124,47 @@ def get_connection_to_mysql( ) return connection + + +def clean_up_connection(connection_config: dict) -> None: + """ + Perform any necessary connection clean up steps after the measuring session. + + :param connection_config: the connection details. + :return: none. + """ + if connection_config["ssh_tunneling"]["use_tunnel"]: + close_ssh_tunnel() + + +def open_ssh_tunnel(connection_config: dict) -> None: + """ + Start an SSH tunnel with the passed details. + + :param connection_config: the connection details. + :return: none. + """ + + print( + f"""Opening up an SSH tunnel to {connection_config["ssh_tunneling"]["ssh_host"]}""" + ) + MySSHTunnel( + ssh_host=connection_config["ssh_tunneling"]["ssh_host"], + ssh_port=connection_config["ssh_tunneling"]["ssh_port"], + ssh_username=connection_config["ssh_tunneling"]["ssh_username"], + ssh_pkey=connection_config["ssh_tunneling"]["path_to_key"], + remote_host=connection_config["host"], + remote_port=connection_config["port"], + ).start() + print("SSH tunnel is now open.") + + +def close_ssh_tunnel(): + """ + Close the SSH tunnel. No details required because a singleton is being used. + + :return: None + """ + print(f"Closing down the SSH tunnel...") + MySSHTunnel().stop() + print("SSH tunnel is now closed.")