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

pg4wp/driver_pgsql.php

@@ -16,6 +16,7 @@ $GLOBALS['pg4wp_result'] = 0;
 $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,23 @@ 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
+        FROM   pg_index i
+        JOIN   pg_attribute a ON a.attrelid = i.indrelid
+                            AND a.attnum = ANY(i.indkey)
+        WHERE  i.indrelid = '$table'::regclass
+        AND    i.indisprimary
+    SQL;
+
+    $result = pg_query($connection, $query);
+
+    $row = pg_fetch_row($result);
+    return $row[0];
+}
+
 /**
  * Performs a query against the database.
  *
@@ -515,6 +533,19 @@ function wpsqli_query(&$connection, $query, $result_mode = 0)
     $GLOBALS['pg4wp_conn'] = $connection;
     $GLOBALS['pg4wp_result'] = $result;
 
+    if (false !== str_post("INSERT INTO")) {
+        $matches = array();
+        preg_match("/^INSERT INTO ([a-z0-9_]+)/i", $query, $matches);
+        $tableName = $matches[1];
+
+        if (false !== str_pos($sql, "RETURNING")) {
+            $primaryKey = $this->get_primary_key_for_table($connection, $tableName);
+            $row = pg_fetch_assoc($result);
+
+            $GLOBALS['pg4wp_ins_id'] = $row[$primaryKey]; 
+        }
+    }
+
     return $result;
 }
 
@@ -1077,9 +1108,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;
 }
 
 /**
@@ -1115,8 +1145,12 @@ function wpsqli_insert_id(&$connection = null)
 
         // 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);";
+    } elseif($GLOBALS['pg4wp_ins_id']) {
+        return $GLOBALS['pg4wp_ins_id']; 
+    } elseif(empty($sql)) {
+        $sql = 'NO QUERY';
+        $data = 0; 
     } else {
-        // PostgreSQL: Using CURRVAL() to get the current value of the sequence.
         // Double quoting is needed to prevent seq from being lowercased automatically
         $sql = "SELECT CURRVAL('\"$seq\"')";
         $res = pg_query($connection, $sql);

pg4wp/rewriters/InsertSQLRewriter.php

@@ -119,6 +119,43 @@ class InsertSQLRewriter extends AbstractSQLRewriter
             $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; 
+    
+            else if($sqlAsArray[$i] == '(' && !$quoteOpened)
+                $parenthesisDepth++;
+    
+            else if($sqlAsArray[$i] == ')' && !$quoteOpened)
+                $parenthesisDepth--;
+    
+            else if($sqlAsArray[$i] == ';' && !$quoteOpened && $parenthesisDepth == 0) 
+                return $i;
+        }
+        return false;
+    }
 }