Ensure all messages are retrieved from the server in the Query Tool. Fixes #3094

web/pgadmin/feature_tests/keyboard_shortcut_test.py

@@ -62,7 +62,9 @@ class KeyboardShortcutFeatureTest(BaseFeatureTest):
             ).key_down(
                 key_combo[2]
             ).key_up(
-                Keys.ALT
+                key_combo[0]
+            ).key_up(
+                key_combo[1]
             ).perform()
 
             print("Executing shortcut: " + self.new_shortcuts[s]['locator'] +

web/pgadmin/tools/sqleditor/tests/test_poll_query_tool.py

@@ -0,0 +1,112 @@
+##########################################################################
+#
+# pgAdmin 4 - PostgreSQL Tools
+#
+# Copyright (C) 2013 - 2018, The pgAdmin Development Team
+# This software is released under the PostgreSQL Licence
+#
+##########################################################################
+
+import json
+
+from pgadmin.browser.server_groups.servers.databases.tests import utils as \
+    database_utils
+from pgadmin.utils.route import BaseTestGenerator
+from regression import parent_node_dict
+from regression.python_test_utils import test_utils as utils
+
+
+class TestPollQueryTool(BaseTestGenerator):
+    """ This class will test the query tool polling. """
+    scenarios = [
+        ('When query tool polling returns messages with result data-set',
+         dict(
+             sql=[
+                 """
+DROP TABLE IF EXISTS test_for_notices;
+
+DO $$
+BEGIN
+    RAISE NOTICE 'Hello, world!';
+END $$;
+
+SELECT 'CHECKING POLLING';
+""",
+                 """
+DO $$
+BEGIN
+    FOR i in 1..1000 LOOP
+        RAISE NOTICE 'Count is %', i;
+    END LOOP;
+END $$;
+
+SELECT 'CHECKING POLLING FOR LONG MESSAGES';
+""",
+                 "SELECT 'CHECKING POLLING WITHOUT MESSAGES';"
+             ],
+             expected_message=['NOTICE:  table "test_for_notices" ' +
+                               """does not exist, skipping
+NOTICE:  Hello, world!
+""",
+                               "\n".join(["NOTICE:  Count is {0}".format(i)
+                                          for i in range(1, 1001)]) + "\n",
+                               None],
+             expected_result=['CHECKING POLLING',
+                              'CHECKING POLLING FOR LONG MESSAGES',
+                              'CHECKING POLLING WITHOUT MESSAGES'],
+             print_messages=['2 NOTICES WITH DATASET',
+                             '1000 NOTICES WITH DATASET',
+                             'NO NOTICE WITH DATASET'
+                             ]
+         ))
+    ]
+
+    def runTest(self):
+        """ This function will check messages return by query tool polling. """
+        database_info = parent_node_dict["database"][-1]
+        self.server_id = database_info["server_id"]
+
+        self.db_id = database_info["db_id"]
+        db_con = database_utils.connect_database(self,
+                                                 utils.SERVER_GROUP,
+                                                 self.server_id,
+                                                 self.db_id)
+        if not db_con["info"] == "Database connected.":
+            raise Exception("Could not connect to the database.")
+
+        # Initialize query tool
+        url = '/datagrid/initialize/query_tool/{0}/{1}/{2}'.format(
+            utils.SERVER_GROUP, self.server_id, self.db_id)
+        response = self.tester.post(url)
+        self.assertEquals(response.status_code, 200)
+
+        response_data = json.loads(response.data.decode('utf-8'))
+        self.trans_id = response_data['data']['gridTransId']
+
+        cnt = 0
+        for s in self.sql:
+            print("Executing and polling with: " + self.print_messages[cnt])
+            # Start query tool transaction
+            url = '/sqleditor/query_tool/start/{0}'.format(self.trans_id)
+            response = self.tester.post(url, data=json.dumps({"sql": s}),
+                                        content_type='html/json')
+
+            self.assertEquals(response.status_code, 200)
+
+            # Query tool polling
+            url = '/sqleditor/poll/{0}'.format(self.trans_id)
+            response = self.tester.get(url)
+            self.assertEquals(response.status_code, 200)
+            response_data = json.loads(response.data.decode('utf-8'))
+
+            # Check the returned messages
+            self.assertEquals(self.expected_message[cnt],
+                              response_data['data']['additional_messages'])
+            # Check the output
+            self.assertEquals(self.expected_result[cnt],
+                              response_data['data']['result'][0][0])
+
+            cnt += 1
+
+        # Disconnect the database
+        database_utils.disconnect_database(self, self.server_id, self.db_id)

web/pgadmin/utils/driver/abstract.py

@@ -168,6 +168,8 @@ class BaseConnection(object):
     ASYNC_WRITE_TIMEOUT = 3
     ASYNC_NOT_CONNECTED = 4
     ASYNC_EXECUTION_ABORTED = 5
+    ASYNC_TIMEOUT = 0.2
+    ASYNC_NOTICE_MAXLENGTH = 100000
 
     @abstractmethod
     def connect(self, **kwargs):

web/pgadmin/utils/driver/psycopg2/__init__.py

@@ -37,6 +37,7 @@ from .cursor import DictCursor
 from .typecast import register_global_typecasters, \
     register_string_typecasters, register_binary_typecasters, \
     register_array_to_string_typecasters, ALL_JSON_TYPES
+from collections import deque
 
 
 if sys.version_info < (3,):
@@ -110,7 +111,7 @@ class Connection(BaseConnection):
       - This method is used to wait for asynchronous connection. This is a
         blocking call.
 
-    * _wait_timeout(conn, time)
+    * _wait_timeout(conn)
       - This method is used to wait for asynchronous connection with timeout.
         This is a non blocking call.
 
@@ -310,6 +311,9 @@ class Connection(BaseConnection):
             )
             return False, msg
 
+        # Overwrite connection notice attr to support
+        # more than 50 notices at a time
+        pg_conn.notices = deque([], self.ASYNC_NOTICE_MAXLENGTH)
         self.conn = pg_conn
         self.wasConnected = True
         try:
@@ -1208,6 +1212,7 @@ Failed to reset the connection to the server due to following error:
             )
             return False, msg
 
+        pg_conn.notices = deque([], self.ASYNC_NOTICE_MAXLENGTH)
         self.conn = pg_conn
         self.__backend_pid = pg_conn.get_backend_pid()
 
@@ -1261,51 +1266,31 @@ Failed to reset the connection to the server due to following error:
 
         Args:
             conn: connection object
-            time: wait time
         """
-
-        state = conn.poll()
-        if state == psycopg2.extensions.POLL_OK:
-            return self.ASYNC_OK
-        elif state == psycopg2.extensions.POLL_WRITE:
-            # Wait for the given time and then check the return status
-            # If three empty lists are returned then the time-out is reached.
-            timeout_status = select.select([], [conn.fileno()], [], 0)
-            if timeout_status == ([], [], []):
-                return self.ASYNC_WRITE_TIMEOUT
-
-            # poll again to check the state if it is still POLL_WRITE
-            # then return ASYNC_WRITE_TIMEOUT else return ASYNC_OK.
+        while 1:
             state = conn.poll()
-            if state == psycopg2.extensions.POLL_WRITE:
-                return self.ASYNC_WRITE_TIMEOUT
-            return self.ASYNC_OK
-        elif state == psycopg2.extensions.POLL_READ:
-            # Wait for the given time and then check the return status
-            # If three empty lists are returned then the time-out is reached.
-            timeout_status = select.select([conn.fileno()], [], [], 0)
-            if timeout_status == ([], [], []):
-                return self.ASYNC_READ_TIMEOUT
-
-            # select.select timeout option works only if we provide
-            #  empty [] [] [] file descriptor in select.select() function
-            # and that also works only on UNIX based system, it do not support
-            # Windows Hence we have wrote our own pooling mechanism to read
-            # data fast each call conn.poll() reads chunks of data from
-            # connection object more we poll more we read data from connection
-            cnt = 0
-            while cnt < 1000:
-                # poll again to check the state if it is still POLL_READ
-                # then return ASYNC_READ_TIMEOUT else return ASYNC_OK.
-                state = conn.poll()
-                if state == psycopg2.extensions.POLL_OK:
-                    return self.ASYNC_OK
-                cnt += 1
-            return self.ASYNC_READ_TIMEOUT
-        else:
-            raise psycopg2.OperationalError(
-                "poll() returned %s from _wait_timeout function" % state
-            )
+            if state == psycopg2.extensions.POLL_OK:
+                return self.ASYNC_OK
+            elif state == psycopg2.extensions.POLL_WRITE:
+                # Wait for the given time and then check the return status
+                # If three empty lists are returned then the time-out is
+                # reached.
+                timeout_status = select.select([], [conn.fileno()], [],
+                                               self.ASYNC_TIMEOUT)
+                if timeout_status == ([], [], []):
+                    return self.ASYNC_WRITE_TIMEOUT
+            elif state == psycopg2.extensions.POLL_READ:
+                # Wait for the given time and then check the return status
+                # If three empty lists are returned then the time-out is
+                # reached.
+                timeout_status = select.select([conn.fileno()], [], [],
+                                               self.ASYNC_TIMEOUT)
+                if timeout_status == ([], [], []):
+                    return self.ASYNC_READ_TIMEOUT
+            else:
+                raise psycopg2.OperationalError(
+                    "poll() returned %s from _wait_timeout function" % state
+                )
 
     def poll(self, formatted_exception_msg=False, no_result=False):
         """
@@ -1347,8 +1332,8 @@ Failed to reset the connection to the server due to following error:
             is_error = True
 
         if self.conn.notices and self.__notices is not None:
-            while self.conn.notices:
-                self.__notices.append(self.conn.notices.pop(0)[:])
+            self.__notices.extend(self.conn.notices)
+            self.conn.notices.clear()
 
         # We also need to fetch notices before we return from function in case
         # of any Exception, To avoid code duplication we will return after