From 90f7aa9e39e744e708a597ff6d1ac83f7c4fcef5 Mon Sep 17 00:00:00 2001
From: Aleksander Machniak <alec@alec.pl>
Date: Tue, 29 Apr 2014 05:34:43 -0400
Subject: [PATCH] Fix varius db_prefix issues (#1489839) - Move DDL script execution code to rcube_db class(es). - Improve prefix replacement code, so index names are also modified

---
 program/lib/Roundcube/rcube_db.php |  168 ++++++++++++++++++++++++++++++++++++++++++++++---------
 1 files changed, 140 insertions(+), 28 deletions(-)

diff --git a/program/lib/Roundcube/rcube_db.php b/program/lib/Roundcube/rcube_db.php
index 2861a91..29f125d 100644
--- a/program/lib/Roundcube/rcube_db.php
+++ b/program/lib/Roundcube/rcube_db.php
@@ -33,6 +33,7 @@
     protected $db_mode;               // Connection mode
     protected $dbh;                   // Connection handle
     protected $dbhs = array();
+    protected $table_connections = array();
 
     protected $db_error     = false;
     protected $db_error_msg = '';
@@ -98,10 +99,15 @@
         $this->db_dsnw  = $db_dsnw;
         $this->db_dsnr  = $db_dsnr;
         $this->db_pconn = $pconn;
-        $this->db_dsnw_noread = rcube::get_instance()->config->get('db_dsnw_noread', false);
 
         $this->db_dsnw_array = self::parse_dsn($db_dsnw);
         $this->db_dsnr_array = self::parse_dsn($db_dsnr);
+
+        $config = rcube::get_instance()->config;
+
+        $this->options['table_prefix']  = $config->get('db_prefix');
+        $this->options['dsnw_noread']   = $config->get('db_dsnw_noread', false);
+        $this->options['table_dsn_map'] = array_map(array($this, 'table_name'), $config->get('db_table_dsn', array()));
     }
 
     /**
@@ -185,8 +191,9 @@
      * Connect to appropriate database depending on the operation
      *
      * @param string $mode Connection mode (r|w)
+     * @param boolean $force Enforce using the given mode
      */
-    public function db_connect($mode)
+    public function db_connect($mode, $force = false)
     {
         // previous connection failed, don't attempt to connect again
         if ($this->conn_failure) {
@@ -201,7 +208,7 @@
         // Already connected
         if ($this->db_connected) {
             // connected to db with the same or "higher" mode (if allowed)
-            if ($this->db_mode == $mode || $this->db_mode == 'w' && !$this->db_dsnw_noread) {
+            if ($this->db_mode == $mode || $this->db_mode == 'w' && !$force && !$this->options['dsnw_noread']) {
                 return;
             }
         }
@@ -215,6 +222,46 @@
         }
 
         $this->conn_failure = !$this->db_connected;
+    }
+
+    /**
+     * Analyze the given SQL statement and select the appropriate connection to use
+     */
+    protected function dsn_select($query)
+    {
+        // no replication
+        if ($this->db_dsnw == $this->db_dsnr) {
+            return 'w';
+        }
+
+        // Read or write ?
+        $mode = preg_match('/^(select|show|set)/i', $query) ? 'r' : 'w';
+
+        // find tables involved in this query
+        if (preg_match_all('/(?:^|\s)(from|update|into|join)\s+'.$this->options['identifier_start'].'?([a-z0-9._]+)'.$this->options['identifier_end'].'?\s+/i', $query, $matches, PREG_SET_ORDER)) {
+            foreach ($matches as $m) {
+                $table = $m[2];
+
+                // always use direct mapping
+                if ($this->options['table_dsn_map'][$table]) {
+                    $mode = $this->options['table_dsn_map'][$table];
+                    break;  // primary table rules
+                }
+                else if ($mode == 'r') {
+                    // connected to db with the same or "higher" mode for this table
+                    $db_mode = $this->table_connections[$table];
+                    if ($db_mode == 'w' && !$this->options['dsnw_noread']) {
+                        $mode = $db_mode;
+                    }
+                }
+            }
+
+            // remember mode chosen (for primary table)
+            $table = $matches[0][2];
+            $this->table_connections[$table] = $mode;
+        }
+
+        return $mode;
     }
 
     /**
@@ -347,12 +394,9 @@
      */
     protected function _query($query, $offset, $numrows, $params)
     {
-        $query = trim($query);
+        $query = ltrim($query);
 
-        // Read or write ?
-        $mode = preg_match('/^(select|show|set)/i', $query) ? 'r' : 'w';
-
-        $this->db_connect($mode);
+        $this->db_connect($this->dsn_select($query), true);
 
         // check connection before proceeding
         if (!$this->is_connected()) {
@@ -363,27 +407,28 @@
             $query = $this->set_limit($query, $numrows, $offset);
         }
 
-        $params = (array) $params;
-
         // Because in Roundcube we mostly use queries that are
         // executed only once, we will not use prepared queries
         $pos = 0;
         $idx = 0;
 
-        while ($pos = strpos($query, '?', $pos)) {
-            if ($query[$pos+1] == '?') {  // skip escaped ?
-                $pos += 2;
-            }
-            else {
-                $val = $this->quote($params[$idx++]);
-                unset($params[$idx-1]);
-                $query = substr_replace($query, $val, $pos, 1);
-                $pos += strlen($val);
+        if (count($params)) {
+            while ($pos = strpos($query, '?', $pos)) {
+                if ($query[$pos+1] == '?') {  // skip escaped '?'
+                    $pos += 2;
+                }
+                else {
+                    $val = $this->quote($params[$idx++]);
+                    unset($params[$idx-1]);
+                    $query = substr_replace($query, $val, $pos, 1);
+                    $pos += strlen($val);
+                }
             }
         }
 
-        // replace escaped ? back to normal
-        $query = rtrim(strtr($query, array('??' => '?')), ';');
+        // replace escaped '?' back to normal, see self::quote()
+        $query = str_replace('??', '?', $query);
+        $query = rtrim($query, " \t\n\r\0\x0B;");
 
         $this->debug($query);
 
@@ -414,7 +459,7 @@
     {
         $error = $this->dbh->errorInfo();
 
-        if (empty($this->options['ignore_key_errors']) || $error[0] != '23000') {
+        if (empty($this->options['ignore_key_errors']) || !in_array($error[0], array('23000', '23505'))) {
             $this->db_error = true;
             $this->db_error_msg = sprintf('[%s] %s', $error[1], $error[2]);
 
@@ -566,7 +611,8 @@
     {
         // get tables if not cached
         if ($this->tables === null) {
-            $q = $this->query('SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES ORDER BY TABLE_NAME');
+            $q = $this->query('SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ? ORDER BY TABLE_NAME',
+                array($this->db_dsnw_array['database']));
 
             if ($q) {
                 $this->tables = $q->fetchAll(PDO::FETCH_COLUMN, 0);
@@ -588,8 +634,8 @@
      */
     public function list_cols($table)
     {
-        $q = $this->query('SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = ?',
-            array($table));
+        $q = $this->query('SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = ? AND TABLE_SCHEMA = ?',
+            array($table, $this->db_dsnw_array['database']));
 
         if ($q) {
             return $q->fetchAll(PDO::FETCH_COLUMN, 0);
@@ -877,10 +923,8 @@
      */
     public function table_name($table)
     {
-        $rcube = rcube::get_instance();
-
         // add prefix to the table name if configured
-        if ($prefix = $rcube->config->get('db_prefix')) {
+        if (($prefix = $this->options['table_prefix']) && strpos($table, $prefix) !== 0) {
             return $prefix . $table;
         }
 
@@ -896,6 +940,17 @@
     public function set_option($name, $value)
     {
         $this->options[$name] = $value;
+    }
+
+    /**
+     * Set DSN connection to be used for the given table
+     *
+     * @param string Table name
+     * @param string DSN connection ('r' or 'w') to be used
+     */
+    public function set_table_dsn($table, $mode)
+    {
+        $this->options['table_dsn_map'][$this->table_name($table)] = $mode;
     }
 
     /**
@@ -1071,4 +1126,61 @@
 
         return $result;
     }
+
+    /**
+     * Execute the given SQL script
+     *
+     * @param string SQL queries to execute
+     *
+     * @return boolen True on success, False on error
+     */
+    public function exec_script($sql)
+    {
+        $sql  = $this->fix_table_names($sql);
+        $buff = '';
+
+        foreach (explode("\n", $sql) as $line) {
+            if (preg_match('/^--/', $line) || trim($line) == '')
+                continue;
+
+            $buff .= $line . "\n";
+            if (preg_match('/(;|^GO)$/', trim($line))) {
+                $this->query($buff);
+                $buff = '';
+                if ($this->db_error) {
+                    break;
+                }
+            }
+        }
+
+        return !$this->db_error;
+    }
+
+    /**
+     * Parse SQL file and fix table names according to table prefix
+     */
+    protected function fix_table_names($sql)
+    {
+        if (!$this->options['table_prefix']) {
+            return $sql;
+        }
+
+        $sql = preg_replace_callback(
+            '/((TABLE|TRUNCATE|(?<!ON )UPDATE|INSERT INTO|FROM'
+            . '| ON(?! (DELETE|UPDATE))|REFERENCES|CONSTRAINT|FOREIGN KEY|INDEX)'
+            . '\s+(IF (NOT )?EXISTS )?[`"]*)([^`"\( \r\n]+)/',
+            array($this, 'fix_table_names_callback'),
+            $sql
+        );
+
+        return $sql;
+    }
+
+    /**
+     * Preg_replace callback for fix_table_names()
+     */
+    protected function fix_table_names_callback($matches)
+    {
+        return $matches[1] . $this->options['table_prefix'] . $matches[count($matches)-1];
+    }
 }

--
Gitblit v1.9.1