Replace calls to seq by adding a RETURNING * to all insert and replace statements

pg4wp/driver_pgsql.php

@@ -16,6 +16,7 @@
 $GLOBALS['pg4wp_numrows_query'] = '';
 $GLOBALS['pg4wp_ins_table'] = '';
 $GLOBALS['pg4wp_ins_field'] = '';
+$GLOBALS['pg4wp_ins_id'] = '';
 $GLOBALS['pg4wp_last_insert'] = '';
 $GLOBALS['pg4wp_connstr'] = '';
 $GLOBALS['pg4wp_conn'] = false;
@@ -465,6 +466,35 @@ function wpsqli_rollback(&$connection, $flags = 0, $name = null)
     pg_query($connection, "ROLLBACK");
 }
 
+function get_primary_key_for_table(&$connection, $table)
+{
+    $query = <<<SQL
+    SELECT a.attname, i.indisprimary 
+        FROM   pg_index i 
+        JOIN   pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
+        WHERE  i.indrelid = '$table'::regclass
+    SQL;
+
+    $result = pg_query($connection, $query);
+    if (!$result) {
+        return null;
+    }
+
+    $firstRow = null;
+    while ($row = pg_fetch_row($result)) {
+        if ($firstRow === null) {
+            $firstRow = $row; // Save the first row in case no match is found
+        }
+
+        if ($row[1] == true) {
+            return $row[0]; // Return the first row where $row[1] == true
+        }
+    }
+
+    // If no row where $row[1] == true was found, return the first row encountered
+    return $firstRow ? $firstRow[0] : null;
+}
+
 /**
  * Performs a query against the database.
  *
@@ -515,6 +545,19 @@ function wpsqli_query(&$connection, $query, $result_mode = 0)
     $GLOBALS['pg4wp_conn'] = $connection;
     $GLOBALS['pg4wp_result'] = $result;
 
+    if (false !== strpos($sql, "INSERT INTO")) {
+        $matches = array();
+        preg_match("/^INSERT INTO\s+`?([a-z0-9_]+)`?/i", $query, $matches);
+        $tableName = $matches[1];
+
+        if (false !== strpos($sql, "RETURNING")) {
+            $primaryKey = get_primary_key_for_table($connection, $tableName);
+            $row = pg_fetch_assoc($result);
+
+            $GLOBALS['pg4wp_ins_id'] = $row[$primaryKey];
+        }
+    }
+
     return $result;
 }
 
@@ -1077,9 +1120,8 @@ function wpsqli_get_primary_sequence_for_table(&$connection, $table)
         }
     }
 
-    // Fallback to default if we don't find a sequence
-    // Note: this will probably fail
-    return $table . '_seq';
+    // we didn't find a sequence for this table.
+    return null;
 }
 
 /**
@@ -1099,24 +1141,15 @@ function wpsqli_insert_id(&$connection = null)
     $data = null;
     $ins_field = $GLOBALS['pg4wp_ins_field'];
     $table = $GLOBALS['pg4wp_ins_table'];
-    $lastq = $GLOBALS['pg4wp_last_insert'];
-    $seq = wpsqli_get_primary_sequence_for_table($connection, $table);
 
-    // Special case when using WP_Import plugin where ID is defined in the query itself.
-    if($table == $wpdb->term_relationships) {
+    if($GLOBALS['pg4wp_ins_id']) {
+        return $GLOBALS['pg4wp_ins_id'];
+    } elseif(empty($sql)) {
         $sql = 'NO QUERY';
         $data = 0;
-    } elseif ('post_author' == $ins_field && false !== strpos($lastq, 'ID')) {
-        // No PostgreSQL specific operation here.
-        $sql = 'ID was in query ';
-        $pattern = '/.+\'(\d+).+$/';
-        preg_match($pattern, $lastq, $matches);
-        $data = $matches[1];
-
-        // PostgreSQL: Setting the value of the sequence based on the latest inserted ID.
-        $GLOBALS['pg4wp_queued_query'] = "SELECT SETVAL('$seq',(SELECT MAX(\"ID\") FROM $table)+1);";
     } else {
-        // PostgreSQL: Using CURRVAL() to get the current value of the sequence.
+        $seq = wpsqli_get_primary_sequence_for_table($connection, $table);
+        $lastq = $GLOBALS['pg4wp_last_insert'];
         // Double quoting is needed to prevent seq from being lowercased automatically
         $sql = "SELECT CURRVAL('\"$seq\"')";
         $res = pg_query($connection, $sql);

pg4wp/rewriters/AlterTableSQLRewriter.php

@@ -36,7 +36,6 @@ public function rewrite(): string
             $sql = $this->rewriteAddIndex($sql);
             return $sql;
         }
-
         if (str_contains($sql, 'CHANGE COLUMN')) {
             $sql = $this->rewriteChangeColumn($sql);
             return $sql;
@@ -65,7 +64,7 @@ public function rewrite(): string
         return $sql;
     }
 
-    private function rewriteAddIndex(string $sql): string 
+    private function rewriteAddIndex(string $sql): string
     {
         $pattern = '/ALTER TABLE\s+(\w+)\s+ADD (UNIQUE |)INDEX\s+([^\s]+)\s+\(((?:[^\(\)]+|\([^\(\)]+\))+)\)/';
 
@@ -74,18 +73,18 @@ private function rewriteAddIndex(string $sql): string
             $unique = $matches[2];
             $index = $matches[3];
             $columns = $matches[4];
-    
+
             // Remove prefix indexing
             // Rarely used and apparently unnecessary for current uses
             $columns = preg_replace('/\([^\)]*\)/', '', $columns);
-    
+
             // Workaround for index name duplicate
             $index = $table . '_' . $index;
-            
+
             // Add backticks around index name and column name, and include IF NOT EXISTS clause
             $sql = "CREATE {$unique}INDEX IF NOT EXISTS `{$index}` ON `{$table}` (`{$columns}`)";
         }
-    
+
         return $sql;
     }
 
@@ -218,15 +217,16 @@ private function rewriteDropPrimaryKey(string $sql): string
         return $sql;
     }
 
-    private function rewrite_numeric_type($sql){
+    private function rewrite_numeric_type($sql)
+    {
         // Numeric types in MySQL which need to be rewritten
         $numeric_types = ["bigint", "int", "integer", "smallint", "mediumint", "tinyint", "double", "decimal"];
         $numeric_types_imploded = implode('|', $numeric_types);
-    
+
         // Prepare regex pattern to match 'type(x)'
         $pattern = "/(" . $numeric_types_imploded . ")\(\d+\)/";
-    
-        // Execute type find & replace 
+
+        // Execute type find & replace
         $sql = preg_replace_callback($pattern, function ($matches) {
             return $matches[1];
         }, $sql);
@@ -260,7 +260,7 @@ private function rewrite_numeric_type($sql){
                 $sql = preg_replace($pattern, 'serial', $sql);
             }
         }
-    
+
         return $sql;
     }
 

pg4wp/rewriters/CreateTableSQLRewriter.php

@@ -89,15 +89,16 @@ public function rewrite(): string
         return $sql;
     }
 
-    private function rewrite_numeric_type($sql){
+    private function rewrite_numeric_type($sql)
+    {
         // Numeric types in MySQL which need to be rewritten
         $numeric_types = ["bigint", "int", "integer", "smallint", "mediumint", "tinyint", "double", "decimal"];
         $numeric_types_imploded = implode('|', $numeric_types);
-    
+
         // Prepare regex pattern to match 'type(x)'
         $pattern = "/(" . $numeric_types_imploded . ")\(\d+\)/";
-    
-        // Execute type find & replace 
+
+        // Execute type find & replace
         $sql = preg_replace_callback($pattern, function ($matches) {
             return $matches[1];
         }, $sql);
@@ -131,7 +132,7 @@ private function rewrite_numeric_type($sql){
                 $sql = preg_replace($pattern, 'serial', $sql);
             }
         }
-    
+
         return $sql;
     }
 

pg4wp/rewriters/InsertSQLRewriter.php

@@ -110,6 +110,42 @@ public function rewrite(): string
             $sql = utf8_encode($sql);
         }
 
+        if(false === strpos($sql, 'RETURNING')) {
+            $end_of_statement = $this->findSemicolon($sql);
+            if ($end_of_statement !== false) {
+                // Create the substrings up to and after the semicolon
+                $sql_before_semicolon = substr($sql, 0, $end_of_statement);
+                $sql_after_semicolon = substr($sql, $end_of_statement, strlen($sql));
+
+                // Splice the SQL string together with 'RETURNING *'
+                $sql = $sql_before_semicolon . ' RETURNING *' . $sql_after_semicolon;
+
+            } else {
+                $sql = $sql .= " RETURNING *";
+            }
+        }
+
         return $sql;
     }
+
+    // finds semicolons that aren't in variables
+    private function findSemicolon($sql)
+    {
+        $quoteOpened = false;
+        $parenthesisDepth = 0;
+
+        $sqlAsArray = str_split($sql);
+        for($i = 0; $i < count($sqlAsArray); $i++) {
+            if(($sqlAsArray[$i] == '"' || $sqlAsArray[$i] == "'") && ($i == 0 || $sqlAsArray[$i - 1] != '\\')) {
+                $quoteOpened = !$quoteOpened;
+            } elseif($sqlAsArray[$i] == '(' && !$quoteOpened) {
+                $parenthesisDepth++;
+            } elseif($sqlAsArray[$i] == ')' && !$quoteOpened) {
+                $parenthesisDepth--;
+            } elseif($sqlAsArray[$i] == ';' && !$quoteOpened && $parenthesisDepth == 0) {
+                return $i;
+            }
+        }
+        return false;
+    }
 }

pg4wp/rewriters/ReplaceIntoSQLRewriter.php

@@ -61,12 +61,12 @@ public function rewrite(): string
             // Extract SQL components
             $tableSection = trim(substr($statement, $insertIndex, $columnsStartIndex - $insertIndex));
             $valuesSection = trim(substr($statement, $valuesIndex, strlen($statement) - $valuesIndex));
-            $columnsSection = trim(substr($statement, $columnsStartIndex, $columnsEndIndex - $columnsStartIndex + 1)); 
+            $columnsSection = trim(substr($statement, $columnsStartIndex, $columnsEndIndex - $columnsStartIndex + 1));
 
             // Extract and clean up column names from the update section
             $updateCols = explode(',', substr($columnsSection, 1, strlen($columnsSection) - 2));
             $updateCols = array_map(function ($col) {
-                return  trim($col); 
+                return  trim($col);
             }, $updateCols);
 
             // Choose a primary key for ON CONFLICT
@@ -91,11 +91,26 @@ public function rewrite(): string
             }
 
             // trim any preceding commas
-            $updateSection = ltrim($updateSection,", ");
+            $updateSection = ltrim($updateSection, ", ");
 
             // Construct the PostgreSQL query
             $postgresSQL = sprintf('%s %s %s ON CONFLICT (%s) DO UPDATE SET %s', $tableSection, $columnsSection, $valuesSection, $primaryKey, $updateSection);
 
+            if(false === strpos($postgresSQL, 'RETURNING')) {
+                $end_of_statement = $this->findSemicolon($postgresSQL);
+                if ($end_of_statement !== false) {
+                    // Create the substrings up to and after the semicolon
+                    $sql_before_semicolon = substr($postgresSQL, 0, $end_of_statement);
+                    $sql_after_semicolon = substr($postgresSQL, $end_of_statement, strlen($postgresSQL));
+
+                    // Splice the SQL string together with 'RETURNING *'
+                    $postgresSQL = $sql_before_semicolon . ' RETURNING *' . $sql_after_semicolon;
+
+                } else {
+                    $postgresSQL = $postgresSQL .= " RETURNING *";
+                }
+            }
+
             // Append to the converted statements list
             $convertedStatements[] = $postgresSQL;
         }
@@ -104,4 +119,25 @@ public function rewrite(): string
 
         return $sql;
     }
+
+      // finds semicolons that aren't in variables
+      private function findSemicolon($sql)
+      {
+          $quoteOpened = false;
+          $parenthesisDepth = 0;
+
+          $sqlAsArray = str_split($sql);
+          for($i = 0; $i < count($sqlAsArray); $i++) {
+              if(($sqlAsArray[$i] == '"' || $sqlAsArray[$i] == "'") && ($i == 0 || $sqlAsArray[$i - 1] != '\\')) {
+                  $quoteOpened = !$quoteOpened;
+              } elseif($sqlAsArray[$i] == '(' && !$quoteOpened) {
+                  $parenthesisDepth++;
+              } elseif($sqlAsArray[$i] == ')' && !$quoteOpened) {
+                  $parenthesisDepth--;
+              } elseif($sqlAsArray[$i] == ';' && !$quoteOpened && $parenthesisDepth == 0) {
+                  return $i;
+              }
+          }
+          return false;
+      }
 }

pg4wp/rewriters/SelectSQLRewriter.php

@@ -37,7 +37,7 @@ public function rewrite(): string
         if(false !== strpos($sql, 'information_schema')) {
             // WP Site Health rewrites
             if (false !== strpos($sql, "SELECT TABLE_NAME AS 'table', TABLE_ROWS AS 'rows', SUM(data_length + index_length)")) {
-                $sql = $this->postgresTableSizeRewrite(); 
+                $sql = $this->postgresTableSizeRewrite();
                 return $sql;
             }
 
@@ -360,7 +360,7 @@ protected function convertToPostgresLimitSyntax($sql)
     }
 
     // This method is specifically to handle should_suggest_persistent_object_cache in wp site health
-    protected function postgresTableSizeRewrite($schema = 'public') 
+    protected function postgresTableSizeRewrite($schema = 'public')
     {
 
         $sql = <<<SQL

pg4wp/rewriters/ShowTableStatusSQLRewriter.php

@@ -5,7 +5,7 @@ class ShowTableStatusSQLRewriter extends AbstractSQLRewriter
     public function rewrite(): string
     {
         $sql = $this->original();
-        return $this->generatePostgresShowTableStatus(); 
+        return $this->generatePostgresShowTableStatus();
     }
 
 

pg4wp/rewriters/ShowVariablesSQLRewriter.php

@@ -38,7 +38,7 @@ public function generatePostgres($sql, $variableName)
         }
 
         if ($variableName == "max_allowed_packet") {
-            // Act like 1GB packet size, in practice this limit doesn't actually exist for postgres, we just want to fool WP 
+            // Act like 1GB packet size, in practice this limit doesn't actually exist for postgres, we just want to fool WP
             return "SELECT '$variableName' AS \"Variable_name\", '1073741824' AS \"Value\";";
         }