fix(dbmanager): wrong set statement timeout

Author: azarz

changelog.md

@@ -1,5 +1,9 @@
 # CHANGELOG
 
+## 3.1.7
+FIXED :
+- wrong database manager logic for set statement timeout
+
 ## 3.1.6
 FIXED :
 - wrong database manager logic for select + batch fetch

r2gg/_database.py

@@ -71,7 +71,10 @@ def connect_working_db(self):
         connect_args = "host=%s dbname=%s user=%s password=%s port=%s" % (host, dbname, user, password, port)
 
         self.logger.info("Connecting to work database")
-        connection = psycopg2.connect(connect_args)
+        if TIMEOUT:
+            connection = psycopg2.connect(connect_args, options=f"-c statement_timeout={TIMEOUT * 1000}")
+        else:
+            connection = psycopg2.connect(connect_args)
         connection.set_client_encoding("UTF8")
 
         return connection
@@ -119,29 +122,30 @@ def execute_select_fetch_multiple(self, query, batchsize=1, show_duration=False)
         """
         self.ensure_connection()
         cursor_name = f"cursor_{int(time.time() * 1000)}"
-        with self._connection.cursor(cursor_factory=DictCursor, name=cursor_name) as cursor:
-            if TIMEOUT:
-                cursor.execute("SET statement_timeout = %s", (1000 * TIMEOUT,))
-            if show_duration:
-                self.logger.info(f"SQL: {query}")
-                st = time.time()
-                cursor.execute(query)
-                self.logger.info(
-                    "Execution ended. Elapsed time : %s seconds.",
-                    time.time() - st
-                )
-            else:
-                cursor.execute(query)
-
-            count = cursor.rowcount
-
-            while True:
-                rows = cursor.fetchmany(batchsize)
-                if not rows:
-                    break
-                if batchsize == 1:
-                    rows = rows.pop()
-                yield rows, count
+        try:
+            with self._connection.cursor(cursor_factory=DictCursor, name=cursor_name) as cursor:
+                if show_duration:
+                    self.logger.info(f"SQL: {query}")
+                    st = time.time()
+                    cursor.execute(query)
+                    self.logger.info(
+                        "Execution ended. Elapsed time : %s seconds.",
+                        time.time() - st
+                    )
+                else:
+                    cursor.execute(query)
+    
+                count = cursor.rowcount
+    
+                while True:
+                    rows = cursor.fetchmany(batchsize)
+                    if not rows:
+                        break
+                    if batchsize == 1:
+                        rows = rows.pop()
+                    yield rows, count
+    finally:
+        self._connection.commit()
 
     # the method below should be used as a generator function otherwise use execute_update
     @database_retry_decorator