From 612f7e734ff29b5a988d6840f0b90cead125bf97 Mon Sep 17 00:00:00 2001
From: Guy Sartorelli <36352093+GuySartorelli@users.noreply.github.com>
Date: Wed, 5 Jul 2023 16:25:58 +1200
Subject: [PATCH] FIX Allow repeated iterations of predicated query result
 (#10857)

---
 src/ORM/Connect/MySQLStatement.php |  6 +++
 tests/php/ORM/DatabaseTest.php     | 68 ++++++++++++++++++++++++++++++
 2 files changed, 74 insertions(+)

diff --git a/src/ORM/Connect/MySQLStatement.php b/src/ORM/Connect/MySQLStatement.php
index 3e95195e9..75d4fca26 100644
--- a/src/ORM/Connect/MySQLStatement.php
+++ b/src/ORM/Connect/MySQLStatement.php
@@ -116,6 +116,12 @@ class MySQLStatement extends Query
             }
             yield $row;
         }
+
+        // Check for the method first since $this->statement isn't strongly typed
+        if (method_exists($this->statement, 'data_seek')) {
+            // Reset so the query can be iterated over again
+            $this->statement->data_seek(0);
+        }
     }
 
     public function numRecords()
diff --git a/tests/php/ORM/DatabaseTest.php b/tests/php/ORM/DatabaseTest.php
index fb01e1085..02ab488a1 100644
--- a/tests/php/ORM/DatabaseTest.php
+++ b/tests/php/ORM/DatabaseTest.php
@@ -312,6 +312,33 @@ class DatabaseTest extends SapphireTest
         $this->assertEquals($inputData, $select->execute()->map());
     }
 
+    /**
+     * Test that repeated abstracted iteration of a query result with predicates returns all records.
+     */
+    public function testRepeatedIterationWithPredicates()
+    {
+        $inputData = ['one', 'two', 'three', 'four'];
+
+        foreach ($inputData as $i => $text) {
+            $x = new MyObject();
+            $x->MyField = $text;
+            $x->MyInt = $i;
+            $x->write();
+        }
+
+        // Note that by including a WHERE statement with predicates
+        // with MySQL the result is in a MySQLStatement object rather than a MySQLQuery object.
+        $select = SQLSelect::create(
+            ['"MyInt"', '"MyField"'],
+            '"DatabaseTest_MyObject"',
+            ['MyInt IN (?,?,?,?,?)' => [0,1,2,3,4]],
+            ['"MyInt"']
+        )->execute();
+
+        $this->assertEquals($inputData, $select->map());
+        $this->assertEquals($inputData, $select->map());
+    }
+
     /**
      * Test that stopping iteration part-way through produces predictable results
      * on a subsequent iteration.
@@ -345,4 +372,45 @@ class DatabaseTest extends SapphireTest
             $i++;
         }
     }
+
+    /**
+     * Test that stopping iteration part-way through produces predictable results even when we're using predicates
+     * on a subsequent iteration.
+     * This test is here to ensure consistency between implementations (e.g. mysql vs postgres, etc)
+     */
+    public function testRepeatedPartialIterationWithPredicates()
+    {
+        $inputData = ['one', 'two', 'three', 'four'];
+
+        foreach ($inputData as $i => $text) {
+            $x = new MyObject();
+            $x->MyField = $text;
+            $x->MyInt = $i;
+            $x->write();
+        }
+
+        // Note that by including a WHERE statement with predicates
+        // with MySQL the result is in a MySQLStatement object rather than a MySQLQuery object.
+        $query = SQLSelect::create(
+            ['"MyInt"', '"MyField"'],
+            '"DatabaseTest_MyObject"',
+            ['MyInt IN (?,?,?,?,?)' => [0,1,2,3,4]],
+            ['"MyInt"']
+        )->execute();
+
+        $i = 0;
+        foreach ($query as $record) {
+            $this->assertEquals($inputData[$i], $record['MyField']);
+            $i++;
+            if ($i > 1) {
+                break;
+            }
+        }
+
+        // Continue from where we left off, since we're using a Generator
+        foreach ($query as $record) {
+            $this->assertEquals($inputData[$i], $record['MyField']);
+            $i++;
+        }
+    }
 }