fix: isWriteType() to recognize CTE; [Postgre] allow beginning whites…

…pace at query start, RETURNING with DELETE
codeigniter4 · Mar 2, 2024 · e70b389 · e70b389
1 parent 9894b92
commit e70b389
Show file tree

Hide file tree

Showing 3 changed files with 137 additions and 8 deletions.
diff --git a/system/Database/BaseConnection.php b/system/Database/BaseConnection.php
@@ -1657,7 +1657,7 @@ public function resetDataCache()
      */
     public function isWriteType($sql): bool
     {
-        return (bool) preg_match('/^\s*"?(SET|INSERT|UPDATE|DELETE|REPLACE|CREATE|DROP|TRUNCATE|LOAD|COPY|ALTER|RENAME|GRANT|REVOKE|LOCK|UNLOCK|REINDEX|MERGE)\s/i', $sql);
+        return (bool) preg_match('/^\s*(WITH\s(\s|.)+(\s|[)]))?"?(SET|INSERT|UPDATE|DELETE|REPLACE|CREATE|DROP|TRUNCATE|LOAD|COPY|ALTER|RENAME|GRANT|REVOKE|LOCK|UNLOCK|REINDEX|MERGE)\s/i', $sql);
     }
 
     /**

diff --git a/system/Database/Postgre/Connection.php b/system/Database/Postgre/Connection.php
@@ -577,7 +577,7 @@ protected function _transRollback(): bool
      */
     public function isWriteType($sql): bool
     {
-        if (preg_match('#^(INSERT|UPDATE).*RETURNING\s.+(\,\s?.+)*$#is', $sql)) {
+        if (preg_match('#^\s*(WITH\s(\s|.)+(\s|[)]))?(INSERT|UPDATE|DELETE).*RETURNING\s.+(\,\s?.+)*$#is', $sql)) {
             return false;
         }
 

diff --git a/tests/system/Database/Live/WriteTypeQueryTest.php b/tests/system/Database/Live/WriteTypeQueryTest.php
@@ -48,11 +48,49 @@ public function testInsert(): void
 
         $this->assertTrue($this->db->isWriteType($sql));
 
-        if ($this->db->DBDriver === 'Postgre') {
-            $sql = "INSERT INTO my_table (col1, col2) VALUES ('Joe', 'Cool') RETURNING id;";
+        $sql = "WITH seqvals AS (SELECT '3' AS seqval)INSERT INTO my_table (col1, col2) SELECT 'Joe', seqval FROM seqvals;";
+
+        $this->assertTrue($this->db->isWriteType($sql));
+
+        $sql = <<<SQL
+            WITH seqvals AS (SELECT '3' AS seqval)
+            INSERT INTO my_table (col1, col2)
+            SELECT 'Joe', seqval
+            FROM seqvals;
+        SQL;
 
-            $this->assertFalse($this->db->isWriteType($sql));
+        $this->assertTrue($this->db->isWriteType($sql));
+
+        $assertionType = 'assertTrue';
+        if ($this->db->DBDriver === 'Postgre') {
+            $assertionType = 'assertFalse';
         }
+
+        $sql = "INSERT INTO my_table (col1, col2) VALUES ('Joe', 'Cool') RETURNING id;";
+
+        $this->$assertionType($this->db->isWriteType($sql));
+
+        $sql = "WITH seqvals AS (SELECT '3' AS seqval)INSERT INTO my_table (col1, col2) SELECT 'Joe', seqval FROM seqvals RETURNING id;";
+
+        $this->$assertionType($this->db->isWriteType($sql));
+
+        $sql = <<<SQL
+            INSERT INTO my_table (col1, col2)
+            VALUES ('Joe', 'Cool')
+            RETURNING id;
+        SQL;
+
+        $this->$assertionType($this->db->isWriteType($sql));
+
+        $sql = <<<SQL
+            WITH seqvals AS (SELECT '3' AS seqval)
+            INSERT INTO my_table (col1, col2)
+            SELECT 'Joe', seqval
+            FROM seqvals
+            RETURNING id;
+        SQL;
+
+        $this->$assertionType($this->db->isWriteType($sql));
     }
 
     public function testUpdate(): void
@@ -63,11 +101,52 @@ public function testUpdate(): void
 
         $this->assertTrue($this->db->isWriteType($sql));
 
-        if ($this->db->DBDriver === 'Postgre') {
-            $sql = "UPDATE my_table SET col1 = 'foo' WHERE id = 2 RETURNING *;";
+        $sql = "WITH seqvals AS (SELECT '3' AS seqval)UPDATE my_table SET col1 = seqval FROM seqvals WHERE id = 2;";
+
+        $this->assertTrue($this->db->isWriteType($sql));
+
+        $sql = <<<SQL
+            WITH seqvals AS (SELECT '3' AS seqval)
+            UPDATE my_table
+            SET col1 = seqval
+            FROM seqvals
+            WHERE id = 2;
+        SQL;
 
-            $this->assertFalse($this->db->isWriteType($sql));
+        $this->assertTrue($this->db->isWriteType($sql));
+
+        $assertionType = 'assertTrue';
+        if ($this->db->DBDriver === 'Postgre') {
+            $assertionType = 'assertFalse';
         }
+
+        $sql = "UPDATE my_table SET col1 = 'foo' WHERE id = 2 RETURNING *;";
+
+        $this->$assertionType($this->db->isWriteType($sql));
+
+        $sql = "WITH seqvals AS (SELECT '3' AS seqval)UPDATE my_table SET col1 = seqval FROM seqvals WHERE id = 2 RETURNING *;";
+
+        $this->$assertionType($this->db->isWriteType($sql));
+
+        $sql = <<<SQL
+            UPDATE my_table
+            SET col1 = 'foo'
+            WHERE id = 2
+            RETURNING *;
+        SQL;
+
+        $this->$assertionType($this->db->isWriteType($sql));
+
+        $sql = <<<SQL
+            WITH seqvals AS (SELECT '3' AS seqval)
+            UPDATE my_table
+            SET col1 = seqval
+            FROM seqvals
+            WHERE id = 2
+            RETURNING *;
+        SQL;
+
+        $this->$assertionType($this->db->isWriteType($sql));
     }
 
     public function testDelete(): void
@@ -76,6 +155,56 @@ public function testDelete(): void
         $sql     = $builder->testMode()->delete(['id' => 1], null, true);
 
         $this->assertTrue($this->db->isWriteType($sql));
+
+        $sql = "DELETE FROM my_table WHERE id = 2;";
+
+        $this->assertTrue($this->db->isWriteType($sql));
+
+        $sql = "WITH seqvals AS (SELECT '3' AS seqval)DELETE FROM my_table JOIN seqvals ON col1 = seqval;";
+
+        $this->assertTrue($this->db->isWriteType($sql));
+
+        $sql = <<<SQL
+            WITH seqvals AS
+            (SELECT '3' AS seqval)
+            DELETE FROM my_table
+            JOIN seqvals
+            ON col1 = seqval;
+        SQL;
+
+        $this->assertTrue($this->db->isWriteType($sql));
+
+        $assertionType = 'assertTrue';
+        if ($this->db->DBDriver === 'Postgre') {
+            $assertionType = 'assertFalse';
+        }
+
+        $sql = "DELETE FROM my_table WHERE id = 2 RETURNING *;";
+
+        $this->$assertionType($this->db->isWriteType($sql));
+
+        $sql = "WITH seqvals AS (SELECT '3' AS seqval)DELETE FROM my_table JOIN seqvals ON col1 = seqval RETURNING *;";
+
+        $this->$assertionType($this->db->isWriteType($sql));
+
+        $sql = <<<SQL
+            DELETE FROM my_table
+            WHERE id = 2
+            RETURNING *;
+        SQL;
+
+        $this->$assertionType($this->db->isWriteType($sql));
+
+        $sql = <<<SQL
+            WITH seqvals AS
+            (SELECT '3' AS seqval)
+            DELETE FROM my_table
+            JOIN seqvals
+            ON col1 = seqval
+            RETURNING *;
+        SQL;
+
+        $this->$assertionType($this->db->isWriteType($sql));
     }
 
     public function testReplace(): void
-Original file line number
+Diff line change
@@ Expand Up / @@ -1657,7 +1657,7 @@ public function resetDataCache() @@
          */
         public function isWriteType($sql): bool
         {
-            return (bool) preg_match('/^\s*"?(SET|INSERT|UPDATE|DELETE|REPLACE|CREATE|DROP|TRUNCATE|LOAD|COPY|ALTER|RENAME|GRANT|REVOKE|LOCK|UNLOCK|REINDEX|MERGE)\s/i', $sql);
+            return (bool) preg_match('/^\s*(WITH\s(\s|.)+(\s|[)]))?"?(SET|INSERT|UPDATE|DELETE|REPLACE|CREATE|DROP|TRUNCATE|LOAD|COPY|ALTER|RENAME|GRANT|REVOKE|LOCK|UNLOCK|REINDEX|MERGE)\s/i', $sql);
         }
         /**
@@ Expand Down @@