<?php
namespace Comuto\Component\Doctrine\DBAL;

use Doctrine\DBAL\Cache\QueryCacheProfile;
use Doctrine\Common\EventManager;
use Doctrine\DBAL\Configuration;
use Doctrine\DBAL\Driver;
use Doctrine\DBAL\Connection as DoctrineConnection;

class Connection extends DoctrineConnection
{
    const MYSQL_CONNECTION_TIMED_WAIT_CODE = 2006;

    public function __construct(array $params, Driver $driver, Configuration $config = null,
            EventManager $eventManager = null)
    {
        parent::__construct($params, $driver, $config, $eventManager);
        $this->_expr = new Query\Expression\ExpressionBuilder($this); /* Get our ExpressionBuilder */
    }

    public function executeQuery($query, array $params = array(), $types = array(), QueryCacheProfile $qcp = null)
    {
        try {
            return parent::executeQuery($query, $params, $types, $qcp);
        } catch (\PDOException $e) {
            if ($e->errorInfo[1] == self::MYSQL_CONNECTION_TIMED_WAIT_CODE) {
                $this->connect();

                return parent::executeQuery($query, $params, $types, $qcp);
            }

            throw $e;
        }
    }

    public function executeUpdate($query, array $params = array(), array $types = array())
    {
        try {
            return parent::executeUpdate($query, $params, $types);
        } catch (\PDOException $e) {
            if ($e->errorInfo[1] == self::MYSQL_CONNECTION_TIMED_WAIT_CODE) {
                $this->connect();

                return parent::executeUpdate($query, $params, $types);
            }

            throw $e;
        }
    }

    public function query()
    {
        try {
            return parent::query();
        } catch (\PDOException $e) {
            if ($e->errorInfo[1] == self::MYSQL_CONNECTION_TIMED_WAIT_CODE) {
                $this->connect();

                return parent::query();
            }

            throw $e;
        }
    }

    public function exec($statement)
    {
        try {
            return parent::exec($statement);
        } catch (\PDOException $e) {
            if ($e->errorInfo[1] == self::MYSQL_CONNECTION_TIMED_WAIT_CODE) {
                $this->connect();

                return parent::exec($statement);
            }

            throw $e;
        }
    }

    public function lastInsertId($seqName = null)
    {
        try {
            return parent::lastInsertId($seqName);
        } catch (\PDOException $e) {
            if ($e->errorInfo[1] == self::MYSQL_CONNECTION_TIMED_WAIT_CODE) {
                $this->connect();

                return parent::lastInsertId($seqName);
            }

            throw $e;
        }
    }

    public function beginTransaction()
    {
        try {
            return parent::beginTransaction();
        } catch (\PDOException $e) {
            if ($e->errorInfo[1] == self::MYSQL_CONNECTION_TIMED_WAIT_CODE) {
                $this->connect();

                return parent::beginTransaction();
            }

            throw $e;
        }
    }

    public function getWrappedConnection()
    {
        try {
            return parent::getWrappedConnection();
        } catch (\PDOException $e) {
            if ($e->errorInfo[1] == self::MYSQL_CONNECTION_TIMED_WAIT_CODE) {
                $this->connect();

                return parent::getWrappedConnection();
            }

            throw $e;
        }
    }

    public function releaseSavepoint($savepoint)
    {
        try {
            return parent::releaseSavepoint($savepoint);
        } catch (\PDOException $e) {
            if ($e->errorInfo[1] == self::MYSQL_CONNECTION_TIMED_WAIT_CODE) {
                $this->connect();

                return parent::releaseSavepoint($savepoint);
            }

            throw $e;
        }
    }

    public function createSavepoint($savepoint)
    {
        try {
            return parent::createSavepoint($savepoint);
        } catch (\PDOException $e) {
            if ($e->errorInfo[1] == self::MYSQL_CONNECTION_TIMED_WAIT_CODE) {
                $this->connect();

                return parent::createSavepoint($savepoint);
            }

            throw $e;
        }
    }

    public function rollbackSavepoint($savepoint)
    {
        try {
            return parent::rollbackSavepoint($savepoint);
        } catch (\PDOException $e) {
            if ($e->errorInfo[1] == self::MYSQL_CONNECTION_TIMED_WAIT_CODE) {
                $this->connect();

                return parent::rollbackSavepoint($savepoint);
            }

            throw $e;
        }
    }

    public function quote($input, $type = null)
    {
        try {
            return parent::quote($input, $type);
        } catch (\PDOException $e) {
            if ($e->errorInfo[1] == self::MYSQL_CONNECTION_TIMED_WAIT_CODE) {
                $this->connect();

                return parent::quote($input, $type);
            }

            throw $e;
        }
    }
}