Override methods for sharding connection

lib/Doctrine/DBAL/Sharding/PoolingShardConnection.php

@@ -19,14 +19,12 @@
 
 namespace Doctrine\DBAL\Sharding;
 
+use Doctrine\Common\EventManager;
+use Doctrine\DBAL\Configuration;
 use Doctrine\DBAL\Connection;
+use Doctrine\DBAL\Driver;
 use Doctrine\DBAL\Event\ConnectionEventArgs;
 use Doctrine\DBAL\Events;
-use Doctrine\DBAL\Driver;
-use Doctrine\DBAL\Configuration;
-
-use Doctrine\Common\EventManager;
-
 use Doctrine\DBAL\Sharding\ShardChoser\ShardChoser;
 
 /**
@@ -128,6 +126,64 @@ class PoolingShardConnection extends Connection
         parent::__construct($params, $driver, $config, $eventManager);
     }
 
+    /**
+     * Get active shard id.
+     * 
+     * @return integer
+     */
+    public function getActiveShardId()
+    {
+        return $this->activeShardId;
+    }
+
+    /**
+     * {@inheritdoc}
+     */
+    public function getParams()
+    {
+        return $this->activeShardId ? $this->connections[$this->activeShardId] : parent::getParams();
+    }
+
+    /**
+     * {@inheritdoc}
+     */
+    public function getHost()
+    {
+        $params = $this->getParams();
+
+        return isset($params['host']) ? $params['host'] : parent::getHost();
+    }
+
+    /**
+     * {@inheritdoc}
+     */
+    public function getPort()
+    {
+        $params = $this->getParams();
+
+        return isset($params['port']) ? $params['port'] : parent::getPort();
+    }
+
+    /**
+     * {@inheritdoc}
+     */
+    public function getUsername()
+    {
+        $params = $this->getParams();
+
+        return isset($params['user']) ? $params['user'] : parent::getUsername();
+    }
+
+    /**
+     * {@inheritdoc}
+     */
+    public function getPassword()
+    {
+        $params = $this->getParams();
+
+        return isset($params['password']) ? $params['password'] : parent::getPassword();
+    }
+
     /**
      * Connects to a given shard.
      *
@@ -151,7 +207,7 @@ class PoolingShardConnection extends Connection
             throw new ShardingException("Cannot switch shard when transaction is active.");
         }
 
-        $this->activeShardId = (int) $shardId;
+        $this->activeShardId = (int)$shardId;
 
         if (isset($this->activeConnections[$this->activeShardId])) {
             $this->_conn = $this->activeConnections[$this->activeShardId];
@@ -211,5 +267,6 @@ class PoolingShardConnection extends Connection
     {
         $this->_conn             = null;
         $this->activeConnections = null;
+        $this->activeShardId     = null;
     }
 }

lib/Doctrine/DBAL/Sharding/PoolingShardManager.php

@@ -19,6 +19,8 @@
 
 namespace Doctrine\DBAL\Sharding;
 
+use Doctrine\DBAL\Sharding\ShardChoser\ShardChoser;
+
 /**
  * Shard Manager for the Connection Pooling Shard Strategy
  *
@@ -27,12 +29,12 @@ namespace Doctrine\DBAL\Sharding;
 class PoolingShardManager implements ShardManager
 {
     /**
-     * @var \Doctrine\DBAL\Sharding\PoolingShardConnection
+     * @var PoolingShardConnection
      */
     private $conn;
 
     /**
-     * @var \Doctrine\DBAL\Sharding\ShardChoser\ShardChoser
+     * @var ShardChoser
      */
     private $choser;
 
@@ -42,7 +44,7 @@ class PoolingShardManager implements ShardManager
     private $currentDistributionValue;
 
     /**
-     * @param \Doctrine\DBAL\Sharding\PoolingShardConnection $conn
+     * @param PoolingShardConnection $conn
      */
     public function __construct(PoolingShardConnection $conn)
     {

tests/Doctrine/Tests/DBAL/Sharding/PoolingShardConnectionTest.php

@@ -20,6 +20,7 @@
 namespace Doctrine\Tests\DBAL\Sharding;
 
 use Doctrine\DBAL\DriverManager;
+use Doctrine\DBAL\Sharding\ShardChoser\MultiTenantShardChoser;
 
 class PoolingShardConnectionTest extends \PHPUnit_Framework_TestCase
 {
@@ -61,7 +62,7 @@ class PoolingShardConnectionTest extends \PHPUnit_Framework_TestCase
     {
         $this->setExpectedException('InvalidArgumentException', "Connection Parameters require 'global' and 'shards' configurations.");
 
-        $conn = DriverManager::getConnection(array(
+        DriverManager::getConnection(array(
             'wrapperClass' => 'Doctrine\DBAL\Sharding\PoolingShardConnection',
             'driver' => 'pdo_sqlite',
             'shards' => array(
@@ -76,7 +77,7 @@ class PoolingShardConnectionTest extends \PHPUnit_Framework_TestCase
     {
         $this->setExpectedException('InvalidArgumentException', "Connection Parameters require 'global' and 'shards' configurations.");
 
-        $conn = DriverManager::getConnection(array(
+        DriverManager::getConnection(array(
             'wrapperClass' => 'Doctrine\DBAL\Sharding\PoolingShardConnection',
             'driver' => 'pdo_sqlite',
             'global' => array('memory' => true),
@@ -88,7 +89,7 @@ class PoolingShardConnectionTest extends \PHPUnit_Framework_TestCase
     {
         $this->setExpectedException('InvalidArgumentException', "Missing Shard Choser configuration 'shardChoser'");
 
-        $conn = DriverManager::getConnection(array(
+        DriverManager::getConnection(array(
             'wrapperClass' => 'Doctrine\DBAL\Sharding\PoolingShardConnection',
             'driver' => 'pdo_sqlite',
             'global' => array('memory' => true),
@@ -103,7 +104,7 @@ class PoolingShardConnectionTest extends \PHPUnit_Framework_TestCase
     {
         $this->setExpectedException('InvalidArgumentException', "The 'shardChoser' configuration is not a valid instance of Doctrine\DBAL\Sharding\ShardChoser\ShardChoser");
 
-        $conn = DriverManager::getConnection(array(
+        DriverManager::getConnection(array(
             'wrapperClass' => 'Doctrine\DBAL\Sharding\PoolingShardConnection',
             'driver' => 'pdo_sqlite',
             'global' => array('memory' => true),
@@ -119,7 +120,7 @@ class PoolingShardConnectionTest extends \PHPUnit_Framework_TestCase
     {
         $this->setExpectedException('InvalidArgumentException', "Shard Id has to be a non-negative number.");
 
-        $conn = DriverManager::getConnection(array(
+        DriverManager::getConnection(array(
             'wrapperClass' => 'Doctrine\DBAL\Sharding\PoolingShardConnection',
             'driver' => 'pdo_sqlite',
             'global' => array('memory' => true),
@@ -134,7 +135,7 @@ class PoolingShardConnectionTest extends \PHPUnit_Framework_TestCase
     {
         $this->setExpectedException('InvalidArgumentException', "Missing 'id' for one configured shard. Please specify a unique shard-id.");
 
-        $conn = DriverManager::getConnection(array(
+        DriverManager::getConnection(array(
             'wrapperClass' => 'Doctrine\DBAL\Sharding\PoolingShardConnection',
             'driver' => 'pdo_sqlite',
             'global' => array('memory' => true),
@@ -149,7 +150,7 @@ class PoolingShardConnectionTest extends \PHPUnit_Framework_TestCase
     {
         $this->setExpectedException('InvalidArgumentException', "Shard 1 is duplicated in the configuration.");
 
-        $conn = DriverManager::getConnection(array(
+        DriverManager::getConnection(array(
             'wrapperClass' => 'Doctrine\DBAL\Sharding\PoolingShardConnection',
             'driver' => 'pdo_sqlite',
             'global' => array('memory' => true),
@@ -178,5 +179,140 @@ class PoolingShardConnectionTest extends \PHPUnit_Framework_TestCase
         $this->setExpectedException('Doctrine\DBAL\Sharding\ShardingException', 'Cannot switch shard when transaction is active.');
         $conn->connect(1);
     }
-}
 
+    public function testGetActiveShardId()
+    {
+        $conn = DriverManager::getConnection(array(
+            'wrapperClass' => 'Doctrine\DBAL\Sharding\PoolingShardConnection',
+            'driver' => 'pdo_sqlite',
+            'global' => array('memory' => true),
+            'shards' => array(
+                array('id' => 1, 'memory' => true),
+            ),
+            'shardChoser' => 'Doctrine\DBAL\Sharding\ShardChoser\MultiTenantShardChoser',
+        ));
+
+        $this->assertNull($conn->getActiveShardId());
+
+        $conn->connect(0);
+        $this->assertEquals(0, $conn->getActiveShardId());
+
+        $conn->connect(1);
+        $this->assertEquals(1, $conn->getActiveShardId());
+
+        $conn->close();
+        $this->assertNull($conn->getActiveShardId());
+    }
+
+    public function testGetParamsOverride()
+    {
+        $conn = DriverManager::getConnection(array(
+            'wrapperClass' => 'Doctrine\DBAL\Sharding\PoolingShardConnection',
+            'driver' => 'pdo_sqlite',
+            'global' => array('memory' => true),
+            'shards' => array(
+                array('id' => 1, 'memory' => true),
+            ),
+            'shardChoser' => 'Doctrine\DBAL\Sharding\ShardChoser\MultiTenantShardChoser',
+        ));
+
+        $this->assertEquals(array(
+            'wrapperClass' => 'Doctrine\DBAL\Sharding\PoolingShardConnection',
+            'driver' => 'pdo_sqlite',
+            'global' => array('memory' => true),
+            'shards' => array(
+                array('id' => 1, 'memory' => true),
+            ),
+            'shardChoser' => new MultiTenantShardChoser(),
+        ), $conn->getParams());
+
+        $conn->connect(1);
+        $this->assertEquals(array(
+            'wrapperClass' => 'Doctrine\DBAL\Sharding\PoolingShardConnection',
+            'driver' => 'pdo_sqlite',
+            'global' => array('memory' => true),
+            'shards' => array(
+                array('id' => 1, 'memory' => true),
+            ),
+            'shardChoser' => new MultiTenantShardChoser(),
+            'id' => 1,
+            'memory' => true,
+        ), $conn->getParams());
+    }
+
+    public function testGetHostOverride()
+    {
+        $conn = DriverManager::getConnection(array(
+            'wrapperClass' => 'Doctrine\DBAL\Sharding\PoolingShardConnection',
+            'driver' => 'pdo_sqlite',
+            'host' => 'localhost',
+            'global' => array('memory' => true),
+            'shards' => array(
+                array('id' => 1, 'memory' => true, 'host' => 'foo'),
+            ),
+            'shardChoser' => 'Doctrine\DBAL\Sharding\ShardChoser\MultiTenantShardChoser',
+        ));
+
+        $this->assertEquals('localhost', $conn->getHost());
+
+        $conn->connect(1);
+        $this->assertEquals('foo', $conn->getHost());
+    }
+
+    public function testGetPortOverride()
+    {
+        $conn = DriverManager::getConnection(array(
+            'wrapperClass' => 'Doctrine\DBAL\Sharding\PoolingShardConnection',
+            'driver' => 'pdo_sqlite',
+            'port' => 3306,
+            'global' => array('memory' => true),
+            'shards' => array(
+                array('id' => 1, 'memory' => true, 'port' => 3307),
+            ),
+            'shardChoser' => 'Doctrine\DBAL\Sharding\ShardChoser\MultiTenantShardChoser',
+        ));
+
+        $this->assertEquals(3306, $conn->getPort());
+
+        $conn->connect(1);
+        $this->assertEquals(3307, $conn->getPort());
+    }
+
+    public function testGetUsernameOverride()
+    {
+        $conn = DriverManager::getConnection(array(
+            'wrapperClass' => 'Doctrine\DBAL\Sharding\PoolingShardConnection',
+            'driver' => 'pdo_sqlite',
+            'user' => 'foo',
+            'global' => array('memory' => true),
+            'shards' => array(
+                array('id' => 1, 'memory' => true, 'user' => 'bar'),
+            ),
+            'shardChoser' => 'Doctrine\DBAL\Sharding\ShardChoser\MultiTenantShardChoser',
+        ));
+
+        $this->assertEquals('foo', $conn->getUsername());
+
+        $conn->connect(1);
+        $this->assertEquals('bar', $conn->getUsername());
+    }
+
+    public function testGetPasswordOverride()
+    {
+        $conn = DriverManager::getConnection(array(
+            'wrapperClass' => 'Doctrine\DBAL\Sharding\PoolingShardConnection',
+            'driver' => 'pdo_sqlite',
+            'password' => 'foo',
+            'global' => array('memory' => true),
+            'shards' => array(
+                array('id' => 1, 'memory' => true, 'password' => 'bar'),
+            ),
+            'shardChoser' => 'Doctrine\DBAL\Sharding\ShardChoser\MultiTenantShardChoser',
+        ));
+
+        $this->assertEquals('foo', $conn->getPassword());
+
+        $conn->connect(1);
+        $this->assertEquals('bar', $conn->getPassword());
+    }
+}