[PBCKP-327] test_ptrack_multiple_segments: try to avoid memory consumption

funny-falcon · funny-falcon · commit 9bcefb2569c3 · 2022-12-02T15:10:22.000+03:00
diff --git a/tests/helpers/ptrack_helpers.py b/tests/helpers/ptrack_helpers.py
@@ -15,6 +15,8 @@
 from time import sleep
 import re
 import json
+from hashlib import md5
+import random
 
 idx_ptrack = {
     't_heap': {
@@ -200,6 +202,32 @@ def kill(self, someone = None):
                 os.kill(self.auxiliary_pids[someone][0], sig)
             self.is_started = False
 
+    def table_checksum(self, table, sort, dbname="postgres"):
+        curname = "cur_"+str(random.randint(0,2**48))
+
+        sum = md5(b"\x01")
+
+        con = self.connect(dbname=dbname)
+
+        con.execute(f"""
+            DECLARE {curname} NO SCROLL CURSOR FOR
+            SELECT t::text FROM {table} as t ORDER BY {sort};
+        """)
+
+        while True:
+            rows = con.execute(f"FETCH FORWARD 10000 FROM {curname}")
+            if not rows:
+                break
+            for row in rows:
+                sum.update(row[0].encode('utf8'))
+                sum.update(b'\x00')
+
+        con.execute(f"CLOSE {curname}; ROLLBACK;")
+
+        con.close()
+        sum.update(b'\x02')
+        return sum.hexdigest()
+
 class ProbackupTest(object):
     # Class attributes
     enterprise = is_enterprise()
diff --git a/tests/ptrack_test.py b/tests/ptrack_test.py
@@ -2039,6 +2039,8 @@ def test_ptrack_multiple_segments(self):
 
         # CREATE TABLE
         node.pgbench_init(scale=100, options=['--tablespace=somedata'])
+        result = node.table_checksum("pgbench_accounts", "aid",
+                                     dbname="postgres")
         # FULL BACKUP
         self.backup_node(backup_dir, 'node', node, options=['--stream'])
 
@@ -2075,7 +2077,8 @@ def test_ptrack_multiple_segments(self):
 
         # GET LOGICAL CONTENT FROM NODE
         # it`s stupid, because hint`s are ignored by ptrack
-        result = node.safe_psql("postgres", "select * from pgbench_accounts")
+        result = node.table_checksum("pgbench_accounts", "aid",
+                                     dbname="postgres")
         # FIRTS PTRACK BACKUP
         self.backup_node(
             backup_dir, 'node', node, backup_type='ptrack', options=['--stream'])
@@ -2108,9 +2111,8 @@ def test_ptrack_multiple_segments(self):
             restored_node, {'port': restored_node.port})
         restored_node.slow_start()
 
-        result_new = restored_node.safe_psql(
-            "postgres",
-            "select * from pgbench_accounts")
+        result_new = restored_node.table_checksum("pgbench_accounts", "aid",
+                                                  dbname="postgres")
 
         # COMPARE RESTORED FILES
         self.assertEqual(result, result_new, 'data is lost')
