Add --shard option to migrations command

Author: vincentchalamon

Command/Helper/DoctrineCommandHelper.php

@@ -15,6 +15,7 @@
 namespace Doctrine\Bundle\MigrationsBundle\Command\Helper;
 
 use Doctrine\Bundle\DoctrineBundle\Command\Proxy\DoctrineCommandHelper as BaseDoctrineCommandHelper;
+use Doctrine\DBAL\Sharding\PoolingShardConnection;
 use Symfony\Bundle\FrameworkBundle\Console\Application;
 use Symfony\Component\Console\Input\InputInterface;
 
@@ -28,13 +29,27 @@ abstract class DoctrineCommandHelper extends BaseDoctrineCommandHelper
 {
     public static function setApplicationHelper(Application $application, InputInterface $input)
     {
-        $doctrine  = $application->getKernel()->getContainer()->get('doctrine');
+        $container = $application->getKernel()->getContainer();
+        $doctrine  = $container->get('doctrine');
         $managerNames = $doctrine->getManagerNames();
 
         if (empty($managerNames)) {
             self::setApplicationConnection($application, $input->getOption('db'));
         } else {
             self::setApplicationEntityManager($application, $input->getOption('em'));
         }
+
+        if ($input->getOption('shard')) {
+            $connection = $application->getHelperSet()->get('db')->getConnection();
+            if (!$connection instanceof PoolingShardConnection) {
+                if (empty($managerNames)) {
+                    throw new \LogicException(sprintf("Connection '%s' must implement shards configuration.", $input->getOption('db')));
+                } else {
+                    throw new \LogicException(sprintf("Connection of EntityManager '%s' must implement shards configuration.", $input->getOption('em')));
+                }
+            }
+
+            $connection->connect($input->getOption('shard'));
+        }
     }
 }

Command/MigrationsDiffDoctrineCommand.php

@@ -14,11 +14,12 @@
 
 namespace Doctrine\Bundle\MigrationsBundle\Command;
 
-use Symfony\Component\Console\Input\InputInterface;
-use Symfony\Component\Console\Output\OutputInterface;
-use Symfony\Component\Console\Input\InputOption;
 use Doctrine\Bundle\DoctrineBundle\Command\Proxy\DoctrineCommandHelper;
 use Doctrine\DBAL\Migrations\Tools\Console\Command\DiffCommand;
+use Doctrine\DBAL\Sharding\PoolingShardConnection;
+use Symfony\Component\Console\Input\InputInterface;
+use Symfony\Component\Console\Input\InputOption;
+use Symfony\Component\Console\Output\OutputInterface;
 
 /**
  * Command for generate migration classes by comparing your current database schema
@@ -36,13 +37,23 @@ protected function configure()
         $this
             ->setName('doctrine:migrations:diff')
             ->addOption('em', null, InputOption::VALUE_OPTIONAL, 'The entity manager to use for this command.')
+            ->addOption('shard', null, InputOption::VALUE_REQUIRED, 'The shard connection to use for this command.')
         ;
     }
 
     public function execute(InputInterface $input, OutputInterface $output)
     {
         DoctrineCommandHelper::setApplicationEntityManager($this->getApplication(), $input->getOption('em'));
 
+        if ($input->getOption('shard')) {
+            $connection = $this->getApplication()->getHelperSet()->get('db')->getConnection();
+            if (!$connection instanceof PoolingShardConnection) {
+                throw new \LogicException(sprintf("Connection of EntityManager '%s' must implements shards configuration.", $input->getOption('em')));
+            }
+
+            $connection->connect($input->getOption('shard'));
+        }
+
         $configuration = $this->getMigrationConfiguration($input, $output);
         DoctrineCommand::configureMigrations($this->getApplication()->getKernel()->getContainer(), $configuration);
 

Command/MigrationsExecuteDoctrineCommand.php

@@ -35,6 +35,7 @@ protected function configure()
             ->setName('doctrine:migrations:execute')
             ->addOption('db', null, InputOption::VALUE_REQUIRED, 'The database connection to use for this command.')
             ->addOption('em', null, InputOption::VALUE_REQUIRED, 'The entity manager to use for this command.')
+            ->addOption('shard', null, InputOption::VALUE_REQUIRED, 'The shard connection to use for this command.')
         ;
     }
 

Command/MigrationsGenerateDoctrineCommand.php

@@ -35,6 +35,7 @@ protected function configure()
             ->setName('doctrine:migrations:generate')
             ->addOption('db', null, InputOption::VALUE_REQUIRED, 'The database connection to use for this command.')
             ->addOption('em', null, InputOption::VALUE_REQUIRED, 'The entity manager to use for this command.')
+            ->addOption('shard', null, InputOption::VALUE_REQUIRED, 'The shard connection to use for this command.')
         ;
     }
 

Command/MigrationsLatestDoctrineCommand.php

@@ -35,6 +35,7 @@ protected function configure()
             ->setName('doctrine:migrations:latest')
             ->addOption('db', null, InputOption::VALUE_REQUIRED, 'The database connection to use for this command.')
             ->addOption('em', null, InputOption::VALUE_REQUIRED, 'The entity manager to use for this command.')
+            ->addOption('shard', null, InputOption::VALUE_REQUIRED, 'The shard connection to use for this command.')
         ;
     }
 

Command/MigrationsMigrateDoctrineCommand.php

@@ -35,6 +35,7 @@ protected function configure()
             ->setName('doctrine:migrations:migrate')
             ->addOption('db', null, InputOption::VALUE_REQUIRED, 'The database connection to use for this command.')
             ->addOption('em', null, InputOption::VALUE_REQUIRED, 'The entity manager to use for this command.')
+            ->addOption('shard', null, InputOption::VALUE_REQUIRED, 'The shard connection to use for this command.')
         ;
     }
 

Command/MigrationsStatusDoctrineCommand.php

@@ -35,6 +35,7 @@ protected function configure()
             ->setName('doctrine:migrations:status')
             ->addOption('db', null, InputOption::VALUE_REQUIRED, 'The database connection to use for this command.')
             ->addOption('em', null, InputOption::VALUE_REQUIRED, 'The entity manager to use for this command.')
+            ->addOption('shard', null, InputOption::VALUE_REQUIRED, 'The shard connection to use for this command.')
         ;
     }
 

Command/MigrationsVersionDoctrineCommand.php

@@ -35,6 +35,7 @@ protected function configure()
             ->setName('doctrine:migrations:version')
             ->addOption('db', null, InputOption::VALUE_REQUIRED, 'The database connection to use for this command.')
             ->addOption('em', null, InputOption::VALUE_REQUIRED, 'The entity manager to use for this command.')
+            ->addOption('shard', null, InputOption::VALUE_REQUIRED, 'The shard connection to use for this command.')
         ;
     }