Merge pull request codeigniter4#8599 from markconnellypro/database-is…

…WriteType fix: isWriteType() to recognize CTE; always excluding RETURNING
najdanovicivan · Mar 7, 2024 · fd28a3f · fd28a3f
2 parents a99787c + feb7f81
commit fd28a3f
Show file tree

Hide file tree

Showing 3 changed files with 275 additions and 30 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|[)]))?"?(SET|INSERT|UPDATE|DELETE|REPLACE|CREATE|DROP|TRUNCATE|LOAD|COPY|ALTER|RENAME|GRANT|REVOKE|LOCK|UNLOCK|REINDEX|MERGE)\s(?!.*\sRETURNING\s)/is', $sql);
     }
 
     /**

diff --git a/system/Database/Postgre/Connection.php b/system/Database/Postgre/Connection.php
@@ -567,20 +567,4 @@ protected function _transRollback(): bool
     {
         return (bool) pg_query($this->connID, 'ROLLBACK');
     }
-
-    /**
-     * Determines if a query is a "write" type.
-     *
-     * Overrides BaseConnection::isWriteType, adding additional read query types.
-     *
-     * @param string $sql
-     */
-    public function isWriteType($sql): bool
-    {
-        if (preg_match('#^(INSERT|UPDATE).*RETURNING\s.+(\,\s?.+)*$#is', $sql)) {
-            return false;
-        }
-
-        return parent::isWriteType($sql);
-    }
 }
diff --git a/tests/system/Database/Live/WriteTypeQueryTest.php b/tests/system/Database/Live/WriteTypeQueryTest.php
@@ -35,7 +35,7 @@ public function testSet(): void
         $this->assertTrue($this->db->isWriteType($sql));
     }
 
-    public function testInsert(): void
+    public function testInsertBuilder(): void
     {
         $builder = $this->db->table('jobs');
 
@@ -47,37 +47,295 @@ public function testInsert(): void
         $sql = $builder->getCompiledInsert();
 
         $this->assertTrue($this->db->isWriteType($sql));
+    }
 
-        if ($this->db->DBDriver === 'Postgre') {
-            $sql = "INSERT INTO my_table (col1, col2) VALUES ('Joe', 'Cool') RETURNING id;";
+    public function testInsertOne(): void
+    {
+        $sql = "INSERT INTO my_table (col1, col2) VALUES ('Joe', 'Cool');";
 
-            $this->assertFalse($this->db->isWriteType($sql));
-        }
+        $this->assertTrue($this->db->isWriteType($sql));
+    }
+
+    public function testInsertMulti(): void
+    {
+        $sql = <<<'SQL'
+                INSERT INTO my_table (col1, col2)
+                VALUES ('Joe', 'Cool');
+            SQL;
+
+        $this->assertTrue($this->db->isWriteType($sql));
+    }
+
+    public function testInsertWithOne(): void
+    {
+        $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));
+    }
+
+    public function testInsertWithOneNoSpace(): void
+    {
+        $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));
     }
 
-    public function testUpdate(): void
+    public function testInsertWithMulti(): void
+    {
+        $sql = <<<'SQL'
+                WITH seqvals AS (SELECT '3' AS seqval)
+                INSERT INTO my_table (col1, col2)
+                SELECT 'Joe', seqval
+                FROM seqvals;
+            SQL;
+
+        $this->assertTrue($this->db->isWriteType($sql));
+    }
+
+    public function testInsertOneReturning(): void
+    {
+        $sql = "INSERT INTO my_table (col1, col2) VALUES ('Joe', 'Cool') RETURNING id;";
+
+        $this->assertFalse($this->db->isWriteType($sql));
+    }
+
+    public function testInsertMultiReturning(): void
+    {
+        $sql = <<<'SQL'
+                INSERT INTO my_table (col1, col2)
+                VALUES ('Joe', 'Cool')
+                RETURNING id;
+            SQL;
+
+        $this->assertFalse($this->db->isWriteType($sql));
+    }
+
+    public function testInsertWithOneReturning(): void
+    {
+        $sql = "WITH seqvals AS (SELECT '3' AS seqval) INSERT INTO my_table (col1, col2) SELECT 'Joe', seqval FROM seqvals RETURNING id;";
+
+        $this->assertFalse($this->db->isWriteType($sql));
+    }
+
+    public function testInsertWithOneReturningNoSpace(): void
+    {
+        $sql = "WITH seqvals AS (SELECT '3' AS seqval)INSERT INTO my_table (col1, col2) SELECT 'Joe', seqval FROM seqvals RETURNING id;";
+
+        $this->assertFalse($this->db->isWriteType($sql));
+    }
+
+    public function testInsertWithMultiReturning(): void
+    {
+        $sql = <<<'SQL'
+                WITH seqvals AS (SELECT '3' AS seqval)
+                INSERT INTO my_table (col1, col2)
+                SELECT 'Joe', seqval
+                FROM seqvals
+                RETURNING id;
+            SQL;
+
+        $this->assertFalse($this->db->isWriteType($sql));
+    }
+
+    public function testUpdateBuilder(): void
     {
         $builder = new BaseBuilder('jobs', $this->db);
         $builder->testMode()->where('id', 1)->update(['name' => 'Programmer'], null, null);
         $sql = $builder->getCompiledInsert();
 
         $this->assertTrue($this->db->isWriteType($sql));
+    }
 
-        if ($this->db->DBDriver === 'Postgre') {
-            $sql = "UPDATE my_table SET col1 = 'foo' WHERE id = 2 RETURNING *;";
+    public function testUpdateOne(): void
+    {
+        $sql = "UPDATE my_table SET col1 = 'foo' WHERE id = 2;";
 
-            $this->assertFalse($this->db->isWriteType($sql));
-        }
+        $this->assertTrue($this->db->isWriteType($sql));
+    }
+
+    public function testUpdateMulti(): void
+    {
+        $sql = <<<'SQL'
+                UPDATE my_table
+                SET col1 = 'foo'
+                WHERE id = 2;
+            SQL;
+
+        $this->assertTrue($this->db->isWriteType($sql));
+    }
+
+    public function testUpdateWithOne(): void
+    {
+        $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));
+    }
+
+    public function testUpdateWithOneNoSpace(): void
+    {
+        $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));
+    }
+
+    public function testUpdateWithMulti(): void
+    {
+        $sql = <<<'SQL'
+                WITH seqvals AS (SELECT '3' AS seqval)
+                UPDATE my_table
+                SET col1 = seqval
+                FROM seqvals
+                WHERE id = 2;
+            SQL;
+
+        $this->assertTrue($this->db->isWriteType($sql));
+    }
+
+    public function testUpdateOneReturning(): void
+    {
+        $sql = "UPDATE my_table SET col1 = 'foo' WHERE id = 2 RETURNING *;";
+
+        $this->assertFalse($this->db->isWriteType($sql));
+    }
+
+    public function testUpdateMultiReturning(): void
+    {
+        $sql = <<<'SQL'
+                UPDATE my_table
+                SET col1 = 'foo'
+                WHERE id = 2
+                RETURNING *;
+            SQL;
+
+        $this->assertFalse($this->db->isWriteType($sql));
+    }
+
+    public function testUpdateWithOneReturning(): void
+    {
+        $sql = "WITH seqvals AS (SELECT '3' AS seqval) UPDATE my_table SET col1 = seqval FROM seqvals WHERE id = 2 RETURNING *;";
+
+        $this->assertFalse($this->db->isWriteType($sql));
+    }
+
+    public function testUpdateWithOneReturningNoSpace(): void
+    {
+        $sql = "WITH seqvals AS (SELECT '3' AS seqval)UPDATE my_table SET col1 = seqval FROM seqvals WHERE id = 2 RETURNING *;";
+
+        $this->assertFalse($this->db->isWriteType($sql));
+    }
+
+    public function testUpdateWithMultiReturning(): void
+    {
+        $sql = <<<'SQL'
+                WITH seqvals AS (SELECT '3' AS seqval)
+                UPDATE my_table
+                SET col1 = seqval
+                FROM seqvals
+                WHERE id = 2
+                RETURNING *;
+            SQL;
+
+        $this->assertFalse($this->db->isWriteType($sql));
     }
 
-    public function testDelete(): void
+    public function testDeleteBuilder(): void
     {
         $builder = $this->db->table('jobs');
         $sql     = $builder->testMode()->delete(['id' => 1], null, true);
 
         $this->assertTrue($this->db->isWriteType($sql));
     }
 
+    public function testDeleteOne(): void
+    {
+        $sql = 'DELETE FROM my_table WHERE id = 2;';
+
+        $this->assertTrue($this->db->isWriteType($sql));
+    }
+
+    public function testDeleteMulti(): void
+    {
+        $sql = <<<'SQL'
+                DELETE FROM my_table
+                WHERE id = 2;
+            SQL;
+
+        $this->assertTrue($this->db->isWriteType($sql));
+    }
+
+    public function testDeleteWithOne(): void
+    {
+        $sql = "WITH seqvals AS (SELECT '3' AS seqval) DELETE FROM my_table JOIN seqvals ON col1 = seqval;";
+
+        $this->assertTrue($this->db->isWriteType($sql));
+    }
+
+    public function testDeleteWithOneNoSpace(): void
+    {
+        $sql = "WITH seqvals AS (SELECT '3' AS seqval)DELETE FROM my_table JOIN seqvals ON col1 = seqval;";
+
+        $this->assertTrue($this->db->isWriteType($sql));
+    }
+
+    public function testDeleteWithMulti(): void
+    {
+        $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));
+    }
+
+    public function testDeleteOneReturning(): void
+    {
+        $sql = 'DELETE FROM my_table WHERE id = 2 RETURNING *;';
+
+        $this->assertFalse($this->db->isWriteType($sql));
+    }
+
+    public function testDeleteMultiReturning(): void
+    {
+        $sql = <<<'SQL'
+                DELETE FROM my_table
+                WHERE id = 2
+                RETURNING *;
+            SQL;
+
+        $this->assertFalse($this->db->isWriteType($sql));
+    }
+
+    public function testDeleteWithOneReturning(): void
+    {
+        $sql = "WITH seqvals AS (SELECT '3' AS seqval) DELETE FROM my_table JOIN seqvals ON col1 = seqval RETURNING *;";
+
+        $this->assertFalse($this->db->isWriteType($sql));
+    }
+
+    public function testDeleteWithOneReturningNoSpace(): void
+    {
+        $sql = "WITH seqvals AS (SELECT '3' AS seqval)DELETE FROM my_table JOIN seqvals ON col1 = seqval RETURNING *;";
+
+        $this->assertFalse($this->db->isWriteType($sql));
+    }
+
+    public function testDeleteWithMultiReturning(): void
+    {
+        $sql = <<<'SQL'
+                WITH seqvals AS
+                (SELECT '3' AS seqval)
+                DELETE FROM my_table
+                JOIN seqvals
+                ON col1 = seqval
+                RETURNING *;
+            SQL;
+
+        $this->assertFalse($this->db->isWriteType($sql));
+    }
+
     public function testReplace(): void
     {
         if (in_array($this->db->DBDriver, ['Postgre', 'SQLSRV'], true)) {
@@ -96,19 +354,22 @@ public function testReplace(): void
         $this->assertTrue($this->db->isWriteType($sql));
     }
 
-    public function testCreate(): void
+    public function testCreateDatabase(): void
     {
         $sql = 'CREATE DATABASE foo';
 
         $this->assertTrue($this->db->isWriteType($sql));
     }
 
-    public function testDrop(): void
+    public function testDropDatabase(): void
     {
         $sql = 'DROP DATABASE foo';
 
         $this->assertTrue($this->db->isWriteType($sql));
+    }
 
+    public function testDropTable(): void
+    {
         $sql = 'DROP TABLE foo';
 
         $this->assertTrue($this->db->isWriteType($sql));