Fix no results to fetch error

Author: vshepard

testgres/connection.py

@@ -102,15 +102,24 @@ def rollback(self):
         return self
 
     def execute(self, query, *args):
-        self.cursor.execute(query, args)
         try:
-            res = self.cursor.fetchall()
-            # pg8000 might return tuples
-            if isinstance(res, tuple):
-                res = [tuple(t) for t in res]
-
-            return res
-        except Exception:
+            self.cursor.execute(query, args)
+            # fetchall works only with select
+            if any(query.strip().lower().startswith(cmd) for cmd in ["select", "show", "fetch"]):
+                res = self.cursor.fetchall()
+                # pg8000 might return tuples
+                if isinstance(res, tuple):
+                    res = [tuple(t) for t in res]
+                return res
+            else:
+                # For INSERT/UPDATE/DELETE return number of rows
+                return self.cursor.rowcount
+
+        except pglib.errors.DuplicateTable as e:
+            print("Error executing query: {}".format(e))
+            return None
+        except QueryException as e:
+            print("Error executing query: {}".format(e))
             return None
 
     def close(self):

testgres/node.py

@@ -1155,28 +1155,38 @@ def poll_query_until(self,
         assert sleep_time > 0
         attempts = 0
         while max_attempts == 0 or attempts < max_attempts:
-            print(f"Pooling {attempts}")
+            print(f"Polling {attempts}")
             try:
-                res = self.execute(dbname=dbname,
-                                   query=query,
-                                   username=username,
-                                   commit=commit)
-
-                if expected is None and res is None:
-                    return    # done
-
-                if res is None:
-                    raise QueryException('Query returned None', query)
-
-                # result set is not empty
-                if len(res):
-                    if len(res[0]) == 0:
-                        raise QueryException('Query returned 0 columns', query)
-                    if res[0][0] == expected:
-                        return    # done
-                # empty result set is considered as None
-                elif expected is None:
-                    return    # done
+                res = self.execute(dbname=dbname, query=query, username=username, commit=commit)
+
+                # Parse res for non-SELECT
+                if not any(query.strip().lower().startswith(cmd) for cmd in ["select", "show", "fetch"]):
+                    if res == expected:
+                        return  # done
+                    elif res is None:
+                        raise QueryException('Query returned None', query)
+                    elif res > -1 and expected is True:
+                        return  # done, assuming non-select query was successful
+                    elif res == -1 and expected is False:
+                        return  # done, assuming non-select query affected no rows
+
+                # Parse res for SELECT
+                else:
+                    if res == expected:
+                        return  # done
+
+                    if res is None:
+                        raise QueryException('Query returned None', query)
+
+                    # result set is not empty
+                    if len(res):
+                        if len(res[0]) == 0:
+                            raise QueryException('Query returned 0 columns', query)
+                        if res[0][0] == expected:
+                            return  # done
+                    # empty result set is considered as None
+                    elif expected is None:
+                        return  # done
 
             except tuple(suppress or []):
                 pass    # we're suppressing them

testgres/plugins/pg_probackup2/setup.py

@@ -4,7 +4,7 @@
     from distutils.core import setup
 
 setup(
-    version='0.0.1',
+    version='0.0.2',
     name='testgres_pg_probackup2',
     packages=['pg_probackup2', 'pg_probackup2.storage'],
     description='Plugin for testgres that manages pg_probackup2',
@@ -14,5 +14,5 @@
     author='Postgres Professional',
     author_email='testgres@postgrespro.ru',
     keywords=['pg_probackup', 'testing', 'testgres'],
-    install_requires=['testgres>=1.9.2']
+    install_requires=['testgres>=1.10.0']
 )