add rewriting for protected column names

PostgreSQL-For-Wordpress · Feb 27, 2024 · b7be1aa · b7be1aa
1 parent 1b6ecce
commit b7be1aa
Show file tree

Hide file tree

Showing 3 changed files with 113 additions and 0 deletions.
diff --git a/pg4wp/rewriters/AlterTableSQLRewriter.php b/pg4wp/rewriters/AlterTableSQLRewriter.php
@@ -30,6 +30,7 @@ public function rewrite(): string
         $sql = $this->original();
 
         $sql = $this->rewrite_numeric_type($sql);
+        $sql = $this->rewrite_columns_with_protected_names($sql);
 
         if (str_contains($sql, 'ADD INDEX') || str_contains($sql, 'ADD UNIQUE INDEX')) {
             $sql = $this->rewriteAddIndex($sql);
@@ -262,4 +263,28 @@ private function rewrite_numeric_type($sql){
 
         return $sql;
     }
+
+    private function rewrite_columns_with_protected_names($sql) 
+    {
+        // Splitting the SQL statement into parts before "(", inside "(", and after ")"
+        if (preg_match('/^(CREATE TABLE IF NOT EXISTS|CREATE TABLE|ALTER TABLE)\s+([^\s]+)\s*\((.*)\)(.*)$/is', $sql, $matches)) {
+            $prefix = $matches[1] . ' ' . $matches[2] . ' (';
+            $columnsAndKeys = $matches[3];
+            $suffix = ')' . $matches[4];
+
+            $regex = '/(?:^|\s*,\s*)(\b(?:timestamp|date|time|default)\b)\s*(?=\s+\w+)/i'; 
+
+            // Callback function to add quotes around protected column names
+            $callback = function($matches) {
+                $whitespace = str_replace($matches[1], "", $matches[0]);
+                return $whitespace . '"' . $matches[1] . '"';
+            };
+
+            // Replace protected column names with quoted versions within columns and keys part
+            $columnsAndKeys = preg_replace_callback($regex, $callback, $columnsAndKeys, 1);
+            return $prefix . $columnsAndKeys . $suffix;
+        }
+
+        return $sql;
+    }
 }
diff --git a/pg4wp/rewriters/CreateTableSQLRewriter.php b/pg4wp/rewriters/CreateTableSQLRewriter.php
@@ -55,6 +55,7 @@ public function rewrite(): string
         );
 
         $sql = $this->rewrite_numeric_type($sql);
+        $sql = $this->rewrite_columns_with_protected_names($sql);
 
         // Support for UNIQUE INDEX creation
         $pattern = '/,\s*(UNIQUE |)KEY\s+(`[^`]+`|\w+)\s+\(((?:[^()]|\([^)]*\))*)\)/';
@@ -133,4 +134,28 @@ private function rewrite_numeric_type($sql){
 
         return $sql;
     }
+
+    private function rewrite_columns_with_protected_names($sql) 
+    {
+        // Splitting the SQL statement into parts before "(", inside "(", and after ")"
+        if (preg_match('/^(CREATE TABLE IF NOT EXISTS|CREATE TABLE|ALTER TABLE)\s+([^\s]+)\s*\((.*)\)(.*)$/is', $sql, $matches)) {
+            $prefix = $matches[1] . ' ' . $matches[2] . ' (';
+            $columnsAndKeys = $matches[3];
+            $suffix = ')' . $matches[4];
+
+            $regex = '/(?:^|\s*,\s*)(\b(?:timestamp|date|time|default)\b)\s*(?=\s+\w+)/i'; 
+
+            // Callback function to add quotes around protected column names
+            $callback = function($matches) {
+                $whitespace = str_replace($matches[1], "", $matches[0]);
+                return $whitespace . '"' . $matches[1] . '"';
+            };
+
+            // Replace protected column names with quoted versions within columns and keys part
+            $columnsAndKeys = preg_replace_callback($regex, $callback, $columnsAndKeys, 1);
+            return $prefix . $columnsAndKeys . $suffix;
+        }
+
+        return $sql;
+    }
 }
diff --git a/tests/rewriteTest.php b/tests/rewriteTest.php
@@ -530,6 +530,69 @@ public function test_it_handles_alter_tables_with_unique_indexes()
         $this->assertSame(trim($expected), trim($postgresql));
     }
 
+    public function test_it_rewrites_protected_column_names()
+    {
+        $sql = <<<SQL
+            CREATE TABLE wp_cmplz_cookiebanners (
+                "ID" int NOT NULL DEFAULT nextval('wp_cmplz_cookiebanners_seq'::text),
+                banner_version int NOT NULL,
+                default int NOT NULL
+            );
+        SQL;
+
+        $expected = <<<SQL
+            CREATE TABLE IF NOT EXISTS wp_cmplz_cookiebanners (
+                 "ID"  int NOT NULL DEFAULT nextval('wp_cmplz_cookiebanners_seq'::text),
+                banner_version int NOT NULL,
+                "default" int NOT NULL
+            );
+        SQL;
+
+        $postgresql = pg4wp_rewrite($sql);
+        $this->assertSame(trim($expected), trim($postgresql));
+    }
+
+    public function test_it_rewrites_advanced_protected_column_names()
+    {
+        $sql = <<<SQL
+            CREATE TABLE wp_statistics_pages (
+                page_id BIGINT(20) NOT NULL AUTO_INCREMENT,
+                uri varchar(190) NOT NULL,
+                type varchar(180) NOT NULL,
+                date date NOT NULL,
+                count int(11) NOT NULL,
+                id int(11) NOT NULL,
+                UNIQUE KEY date_2 (date,uri),
+                KEY url (uri),
+                KEY date (date),
+                KEY id (id),
+                KEY `uri` (`uri`,`count`,`id`),
+                PRIMARY KEY (`page_id`)
+            ) DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_520_ci
+        SQL;
+
+        $expected = <<<SQL
+           CREATE TABLE IF NOT EXISTS wp_statistics_pages (
+                page_id bigserial,
+                uri varchar(190) NOT NULL,
+                type varchar(180) NOT NULL,
+                "date" date NOT NULL,
+                count int NOT NULL,
+                id int NOT NULL,
+                PRIMARY KEY (page_id)
+            );
+        CREATE UNIQUE INDEX IF NOT EXISTS wp_statistics_pages_date_2 ON wp_statistics_pages (date,uri);
+        CREATE INDEX IF NOT EXISTS wp_statistics_pages_url ON wp_statistics_pages (uri);
+        CREATE INDEX IF NOT EXISTS wp_statistics_pages_date ON wp_statistics_pages (date);
+        CREATE INDEX IF NOT EXISTS wp_statistics_pages_id ON wp_statistics_pages (id);
+        CREATE INDEX IF NOT EXISTS wp_statistics_pages_uri ON wp_statistics_pages (uri,count,id);
+        SQL;
+
+        $postgresql = pg4wp_rewrite($sql);
+        $this->assertSame(trim($expected), trim($postgresql));
+    }
+
+
     public function test_it_doesnt_remove_single_quotes() 
     {
         $sql = <<<SQL