Add support for SSH tunneled connections. Fixes #1447

web/pgadmin/tools/datagrid/__init__.py

@@ -27,7 +27,7 @@ from config import PG_DEFAULT_DRIVER
 from pgadmin.utils.preferences import Preferences
 from pgadmin.model import Server
 from pgadmin.utils.driver import get_driver
-from pgadmin.utils.exception import ConnectionLost
+from pgadmin.utils.exception import ConnectionLost, SSHTunnelConnectionLost
 from pgadmin.tools.sqleditor.utils.query_tool_preferences import \
     get_query_tool_keyboard_shortcuts, get_text_representation_of_shortcut
 
@@ -135,7 +135,7 @@ def initialize_datagrid(cmd_type, obj_type, sgid, sid, did, obj_id):
                                   auto_reconnect=False,
                                   use_binary_placeholder=True,
                                   array_to_string=True)
-    except ConnectionLost as e:
+    except (ConnectionLost, SSHTunnelConnectionLost) as e:
         raise
     except Exception as e:
         app.logger.error(e)
@@ -363,7 +363,7 @@ def initialize_query_tool(sgid, sid, did=None):
                                   array_to_string=True)
         if connect:
             conn.connect()
-    except ConnectionLost as e:
+    except (ConnectionLost, SSHTunnelConnectionLost) as e:
         raise
     except Exception as e:
         app.logger.error(e)

web/pgadmin/tools/sqleditor/__init__.py

@@ -34,7 +34,7 @@ from pgadmin.utils.ajax import make_json_response, bad_request, \
     success_return, internal_server_error, unauthorized
 from pgadmin.utils.driver import get_driver
 from pgadmin.utils.menu import MenuItem
-from pgadmin.utils.exception import ConnectionLost
+from pgadmin.utils.exception import ConnectionLost, SSHTunnelConnectionLost
 from pgadmin.utils.sqlautocomplete.autocomplete import SQLAutoComplete
 from pgadmin.tools.sqleditor.utils.query_tool_preferences import \
     RegisterQueryToolPreferences
@@ -166,7 +166,7 @@ def check_transaction_status(trans_id):
             use_binary_placeholder=True,
             array_to_string=True
         )
-    except ConnectionLost as e:
+    except (ConnectionLost, SSHTunnelConnectionLost) as e:
         raise
     except Exception as e:
         current_app.logger.error(e)
@@ -212,7 +212,7 @@ def start_view_data(trans_id):
         manager = get_driver(PG_DEFAULT_DRIVER).connection_manager(
             trans_obj.sid)
         default_conn = manager.connection(did=trans_obj.did)
-    except ConnectionLost as e:
+    except (ConnectionLost, SSHTunnelConnectionLost) as e:
         raise
     except Exception as e:
         current_app.logger.error(e)
@@ -261,7 +261,7 @@ def start_view_data(trans_id):
         # Execute sql asynchronously
         try:
             status, result = conn.execute_async(sql)
-        except ConnectionLost as e:
+        except (ConnectionLost, SSHTunnelConnectionLost) as e:
             raise
     else:
         status = False

web/pgadmin/tools/sqleditor/utils/start_running_query.py

@@ -25,7 +25,7 @@ from pgadmin.tools.sqleditor.utils.update_session_grid_transaction import \
     update_session_grid_transaction
 from pgadmin.utils.ajax import make_json_response, internal_server_error
 from pgadmin.utils.driver import get_driver
-from pgadmin.utils.exception import ConnectionLost
+from pgadmin.utils.exception import ConnectionLost, SSHTunnelConnectionLost
 
 
 class StartRunningQuery:
@@ -61,7 +61,7 @@ class StartRunningQuery:
                                           auto_reconnect=False,
                                           use_binary_placeholder=True,
                                           array_to_string=True)
-            except ConnectionLost:
+            except (ConnectionLost, SSHTunnelConnectionLost):
                 raise
             except Exception as e:
                 self.logger.error(e)
@@ -127,7 +127,7 @@ class StartRunningQuery:
         # and formatted_error is True.
         try:
             status, result = conn.execute_async(sql)
-        except ConnectionLost:
+        except (ConnectionLost, SSHTunnelConnectionLost):
             raise
 
         # If the transaction aborted for some reason and

web/pgadmin/tools/sqleditor/utils/tests/test_start_running_query.py

@@ -12,7 +12,7 @@ from flask import Response
 import simplejson as json
 
 from pgadmin.tools.sqleditor.utils.start_running_query import StartRunningQuery
-from pgadmin.utils.exception import ConnectionLost
+from pgadmin.utils.exception import ConnectionLost, SSHTunnelConnectionLost
 from pgadmin.utils.route import BaseTestGenerator
 
 if sys.version_info < (3, 3):
@@ -176,6 +176,35 @@ class StartRunningQueryTest(BaseTestGenerator):
              is_rollback_required=False,
              apply_explain_plan_wrapper_if_needed_return_value='some sql',
 
+             expect_make_json_response_to_have_been_called_with=None,
+             expect_internal_server_error_called_with=None,
+             expected_logger_error=None,
+             expect_execute_void_called_with='some sql',
+         )),
+        ('When SSHTunnelConnectionLost happens while retrieving the '
+         'database connection, '
+         'it returns an error',
+         dict(
+             function_parameters=dict(
+                 sql=dict(sql='some sql', explain_plan=None),
+                 trans_id=123,
+                 http_session=dict(gridData={'123': dict(command_obj='')})
+             ),
+             pickle_load_return=MagicMock(
+                 conn_id=1,
+                 update_fetched_row_cnt=MagicMock()
+             ),
+             get_driver_exception=False,
+             get_connection_lost_exception=False,
+             manager_connection_exception=SSHTunnelConnectionLost('1.1.1.1'),
+
+             is_connected_to_server=False,
+             connection_connect_return=None,
+             execute_async_return_value=None,
+             is_begin_required=False,
+             is_rollback_required=False,
+             apply_explain_plan_wrapper_if_needed_return_value='some sql',
+
              expect_make_json_response_to_have_been_called_with=None,
              expect_internal_server_error_called_with=None,
              expected_logger_error=None,