DB.php

<?php

function DB($query, $args = NULL, $return = false) {
    $stmt = DB::prepare($query);
    if ($args !== NULL && !is_array($args) && func_num_args() == 2) {
        //support DB(query,return)
        $return = $args;
        $args = array();
    }
    if (!is_null($args) && is_array($args)) {
        $time = microtime(true);
        $stmt->execute($args);
        //echo "\n<br>=====\n<br/>".$query . " - " . (microtime(true) - $time) . "\n<br>====\n<br>";
        if (true) {
            if ($return === 1) {
                return $stmt->fetch();
            } else if ($return === 'all') {
                return $stmt->fetchAll();
            } else if ($return == true) {
                return $stmt->fetchColumn();
            }
        }
    }
    return $stmt;
}

/**
 * @uses PDOStatement
 */
class PDOStatementWrapper {

    private $stmt, $rebuild, $hasexecuted = false;

    public function __construct($stmt, $rebuild) {
        $this->rebuild = $rebuild;
        $stmt->setFetchMode(PDO::FETCH_ASSOC);
        $this->stmt = $stmt;
    }

    public function __call($func, $args) {
        if (!$this->hasexecuted && in_array($func, array('fetch', 'fetchAll', 'fetchColumn'))) {
            $this->execute(array());
        }
        $retry = 0;
        do {
            try {

                return call_user_func_array(array($this->stmt, $func), $args);
            } catch (Exception $e) {
                if (false !== strpos($e->getMessage(), '2006 MySQL server')) {
                    DB::reconnect();
                    $this->stmt = call_user_func_array(array(DB::getInstance($this->rebuild['k']), $this->rebuild['f']), $this->rebuild['a']);
                    $this->stmt->setFetchMode(PDO::FETCH_ASSOC);
                } else {
                    throw $e;
                }
            }
        } while ($retry++ < 2);
        return false;
    }

    public function lastInsertId() {
        return DB::lastInsertId();
    }

    public function execute() {
        $retry = 0;
        do {
            try {
                if (call_user_func_array(array($this->stmt, 'execute'), func_get_args())) {
                    $this->hasexecuted = true;
                    return $this->stmt;
                } else {
                    return false;
                }
            } catch (Exception $e) {
                if (false !== strpos($e->getMessage(), '2006 MySQL server')) {
                    DB::reconnect();
                    $this->stmt = call_user_func_array(array(DB::getInstance($this->rebuild['k']), $this->rebuild['f']), $this->rebuild['a']);
                    $this->stmt->setFetchMode(PDO::FETCH_ASSOC);
                } else {
                    throw $e;
                }
            }
        } while ($retry++ < 2);
        return false;
    }

    public function __invoke() {
        return call_user_func_array(array($this, 'execute'), func_get_args());
    }

}

class DB {
    const WRITE_QUERIES = '/^\s*(UPDATE|INSERT|ALTER|TRUNCATE|DELETE)/i';

    protected static
    $dbInstance = array('read' => NULL, 'write' => NULL),
    $dbDsn = array('read' => NULL, 'write' => NULL),
    $lastDb = NULL;

    /**
     * Private constructor to prevent instantiation of this class.
     */
    private function __construct() {

    }

    /**
     * Initializes the connection info for the Read server and the Write Server
     * @param array $readDsn
     * @param array $writeDsn
     * @return
     */
    public static function initDB($readDsn, $writeDsn = false) {
        if (!$writeDsn)
            $writeDsn = $readDsn;
        self::$dbDsn['read'] = $readDsn;
        self::$dbDsn['write'] = $writeDsn;
    }

    /**
     * utility func to get a repeated string of ?'s for inserts
     * @param object $items
     * @return
     */
    public static function get_placeholder_string($items) {
        return implode(', ', array_fill(0, count((array) $items), '?'));
    }

    /**
     * Same thing as PDO's exec just wrapped to determine READ server vs WRITE server
     * @param object $query
     * @return
     */
    static public function exec($query) {
        $retry = 0;
        do {
            try {
                if (preg_match('/^\s*(UPDATE|INSERT|ALTER|TRUNCATE|DELETE)/i', $query, $match)) {
                    $key = false;
                } else {
                    $key = true;
                }
                self::$lastDb = self::getInstance($key);
                return call_user_func_array(array(self::$lastDb, 'exec'), func_get_args());
            } catch (Exception $e) {
                if (false !== strpos($e->getMessage(), '2006 MySQL server')) {
                    DB::reconnect();
                } else {
                    throw $e;
                }
            }
        } while ($retry++ < 2);
        return false;
    }

    /**
     * Same thing as PDO's query just wrapped to determine READ server vs WRITE server, and
     * wraps the PDOStatement with our Utility wrapper to auto set FETCH_ASSOC and make invokeable.
     * @param object $query
     * @return
     */
    static public function query($query) {
        $retry = 0;
        do {
            try {
                if (preg_match('/^\s*(UPDATE|INSERT|ALTER|TRUNCATE|DELETE)/i', $query, $match)) {
                    $key = false;
                } else {
                    $key = true;
                }
                self::$lastDb = self::getInstance($key);
                return new PDOStatementWrapper(
                        call_user_func_array(array(self::$lastDb, 'query'), func_get_args()), //call
                        array('k' => $key, 'f' => 'query', 'a' => func_get_args())
                );
            } catch (Exception $e) {
                if (false !== strpos($e->getMessage(), '2006 MySQL server')) {
                    self::reconnect();
                } else {
                    throw $e;
                }
            }
        } while ($retry++ < 2);
        return false;
    }

    /**
     * Same thing as PDO's prepare just wrapped to determine READ server vs WRITE server, and
     * wraps the PDOStatement with our Utility wrapper to auto set FETCH_ASSOC and make invokeable.
     * @param object $query
     * @param object $driver_options [optional]
     * @return
     */
    static public function prepare($query, $driver_options = array()) {
        $retry = 0;
        do {
            try {
                if (preg_match('/^\s*(UPDATE|INSERT|ALTER|TRUNCATE|DELETE)/i', $query, $match)) {
                    $key = false;
                } else {
                    $key = true;
                }
                self::$lastDb = self::getInstance($key);
                return new PDOStatementWrapper(
                        call_user_func_array(array(self::$lastDb, 'prepare'), func_get_args()), //call
                        array('k' => $key, 'f' => 'prepare', 'a' => func_get_args()) //rebuild
                );
            } catch (Exception $e) {
                if (false !== strpos($e->getMessage(), '2006 MySQL server')) {
                    self::reconnect();
                } else {
                    throw $e;
                }
            }
        } while ($retry++ < 2);
        return false;
    }

    /**
     * Gets an instance of the DB connection
     * @param object $read [optional]
     * @return
     */
    public static function getInstance($read = true) {
        $key = 'read';
        if (!$read)
            $key = 'write';
        if (is_null(self::$dbInstance[$key])) {
            try {
                $dsn = self::$dbDsn[$key];
                self::$dbInstance[$key] = new PDO($dsn['driver'] . ':host=' . $dsn['host'] . ';dbname=' . $dsn['database'], $dsn['username'], $dsn['password']);
            } catch (Exception $e) {
                die('Could not connect to database.');
            }

            self::$dbInstance[$key]->setAttribute(PDO::ATTR_ERRMODE, PDO::ERRMODE_EXCEPTION);
        }

        return self::$dbInstance[$key];
    }

    public static function reconnect() {
        echo "DB Reconnected\n";
        self::$dbInstance = array('read' => NULL, 'write' => NULL);
        self::$lastDb = self::$dbInstance;
    }

    /**
     * Forwarder wrapper to forward all calls to DB:: to the READ server.
     * @param object $method
     * @param object $args
     * @return
     */
    final public static function __callStatic($method, $args) {
        $retry = 0;
        do {
            try {
                if (!self::$lastDb)
                    self::$lastDb = self::getInstance(true);
                return call_user_func_array(array(self::$lastDb, $method), $args);
            } catch (Exception $e) {
                if (false !== strpos($e->getMessage(), '2006 MySQL server')) {
                    self::reconnect();
                } else {
                    throw $e;
                }
            }
        } while ($retry++ < 2);
        return false;
    }

}