Extended PDO class that detects dropped connections and reconnects

RobustPDO.php

<?php

class RobustPDO extends PDO
{
    /** Call setAttribute to set the session wait_timeout value */
    const ATTR_MYSQL_TIMEOUT = 100;

    /** @var array */
    protected $config = [];
    /** @var bool For lazy connection tracking */
    protected $_connected = false;
    /** @var array Cached attributes for reconnection */
    private $attributes = [];

    /**
     * Create a new PDO object.
     * Does not connect to the database until needed.
     *
     * @param string $dsn     The Data Source Name, or DSN, contains the information required to connect to the database.
     * @param string $user    [optional] The user name for the DSN string. This parameter is optional for some PDO drivers.
     * @param string $pass    [optional] The password for the DSN string. This parameter is optional for some PDO drivers.
     * @param string $options [optional] A key=>value array of driver-specific connection options.
     */
    public function __construct($dsn, $user = null, $pass = null, $options = null) {
        //Save connection details for later
        $this->config = array(
            'dsn'     => $dsn,
            'user'    => $user,
            'pass'    => $pass,
            'options' => $options
        );

        if (!is_array($this->config['options'])) { $this->config['options'] = []; }
        // Throw exceptions when there's an error so we can catch them
        $this->config['options'][PDO::ATTR_ERRMODE] = PDO::ERRMODE_EXCEPTION;
        $this->config['options'][PDO::ATTR_DEFAULT_FETCH_MODE] = PDO::FETCH_OBJ;
        if (isset($this->config['options'][EosPDO::ATTR_MYSQL_TIMEOUT])) {
            $this->attributes[EosPDO::ATTR_MYSQL_TIMEOUT] = $this->config['options'][EosPDO::ATTR_MYSQL_TIMEOUT];
        }
    }

    /**
     * Verifies that the PDO connection is still active. If not, it reconnects.
     */
    public function reconnect() {
        parent::__construct($this->config['dsn'], $this->config['user'], $this->config['pass'], $this->config['options']);
        
        // Reapply attributes to the new connection
        foreach ($this->attributes as $attr => $value) {
            $this->_setAttribute($attr, $value);
        }
        $this->_connected = true;
    }

    /**
     * Try to call the function on the parent inside a try/catch to detect a disconnect
     * 
     * @throws PDOException
     */
    public function __call($name, $arguments) {
        if (!$this->_connected) $this->reconnect();
        try {
            return call_user_func_array(['parent', $name], $arguments);
        } catch (PDOException $e) {
            if (static::hasGoneAway($e)) {
                $this->reconnect();
                return call_user_func_array(['parent', $name], $arguments);
            } else {
                throw $e;
            }
        }
    }

    /**
     * {@InheritDoc}
     */
    public function beginTransaction() {
        return $this->__call(__FUNCTION__, func_get_args());
    }

    /**
     * {@InheritDoc}
     */
    public function query($statement) {
        return $this->__call(__FUNCTION__, func_get_args());
    }

    /**
     * {@InheritDoc}
     */
    public function exec($statement) {
        return $this->__call(__FUNCTION__, func_get_args());
    }

    /**
     * {@InheritDoc}
     */
    public function getAttribute($statement) {
        return $this->__call(__FUNCTION__, func_get_args());
    }

    /**
     * {@InheritDoc}
     */
    public function prepare($statement, $driver_options = null) {
        if (!$this->_connected) $this->reconnect();
        if (is_null($driver_options)) $driver_options = [];
        try {
            return parent::prepare($statement, $driver_options);
        } catch (PDOException $e) {
            if (static::hasGoneAway($e)) {
                $this->reconnect();
                return parent::prepare($statement, $driver_options);
            } else {
                throw $e;
            }
        }
    }

    /**
     * {@InheritDoc}
     * 
     * Caches setAttribute calls to be reapplied on reconnect
     */
    public function setAttribute($attribute, $value) {
        $this->attributes[$attribute] = $value;
        if ($this->_connected) {
            try {
                $this->_setAttribute($attribute, $value);
            } catch (PDOException $e) {
                if (static::hasGoneAway($e)) {
                    $this->reconnect();
                } else {
                    throw $e;
                }
            }
        }
        return true;
    }

    /**
     * Executes a prepared statement inside a try/catch to watch for server disconnections
     *
     * Use to detect connection drops between prepare and PDOStatement->execute
     * Usage:
     * <code>
     * $st = $this->robustPDO->prepare('SELECT :one');
     * if ($this->robustPDO->tryExecuteStatement($st, ['one'=>1])) { print_r($st->fetchAll()); }
     * </code>
     *
     * @param PDOStatement $st                PDOStatement to execute
     * @param null|array   $input_parameters  An array of values with as many elements as there are bound parameters in the SQL statement
     * @param bool         $recursion         If TRUE, don't retry the operation
     * @throws PDOException If the statement throws an exception that's not a disconnect
     * @return bool Returns TRUE on success or FALSE on failure.
     */
    public function tryExecuteStatement(PDOStatement &$st, $input_parameters = null, $recursion = false) {
        try {
            if (is_array($input_parameters))
                return $st->execute($input_parameters);
            else
                return $st->execute();
        } catch (PDOException $e) {
            if (!$recursion && static::hasGoneAway($e)) {
                $this->reconnect();
                $st = $this->prepare($st->queryString);
                return $this->tryExecuteStatement($st, $input_parameters, true);
            }

            throw $e;
        }
    }

    /**
     * Check if a PDOException is a server disconnection or not
     *
     * @param PDOException $e
     * @return bool Returns TRUE if the PDOException is for "server has gone away" or FALSE for another error
     */
    public static function hasGoneAway(PDOException $e) {
        return ($e->getCode() == 'HY000' && stristr($e->getMessage(), 'server has gone away'));
    }

    /**
     * Ping the PDO connection to keep it alive by sending a SELECT 1
     *
     * @return bool
     */
    public function ping() {
        return (1 === intval($this->query('SELECT 1')->fetchColumn(0)));
    }

    /**
     * Sets a connection attribute. Adds support for additional custom attributes from this class.
     *
     * This calls setAttribute right away and requires that the PDO connection be open.
     *
     * @param int   $attribute
     * @param mixed $value
     * @return bool Returns TRUE on success or FALSE on failure
     */
    protected function _setAttribute($attribute, $value) {
        if ($attribute === EosPDO::ATTR_MYSQL_TIMEOUT) {
            parent::exec("SET SESSION wait_timeout=" . intval($value));
            return true;
        } else {
            return parent::setAttribute($attribute, $value);
        }
    }

}