[DoctrineBridge] doctrine connection listener for long running runtime

alli83 · alli83 · commit 7791eb2f2a03 · 2024-03-12T15:19:36.000+11:00
diff --git a/src/Symfony/Bridge/Doctrine/CheckConnection.php b/src/Symfony/Bridge/Doctrine/CheckConnection.php
@@ -0,0 +1,28 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <fabien@symfony.com>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\Bridge\Doctrine;
+
+use Doctrine\DBAL\Connection;
+use Doctrine\DBAL\Exception as DBALException;
+
+class CheckConnection
+{
+    public static function checkConnection(Connection $connection) {
+        try {
+            $connection->executeQuery($connection->getDatabasePlatform()->getDummySelectSQL());
+        } catch (DBALException) {
+            $connection->close();
+            // Attempt to reestablish the lazy connection by sending another query.
+            $connection->executeQuery($connection->getDatabasePlatform()->getDummySelectSQL());
+        }
+    }
+}
diff --git a/src/Symfony/Bridge/Doctrine/Messenger/DoctrinePingConnectionMiddleware.php b/src/Symfony/Bridge/Doctrine/Messenger/DoctrinePingConnectionMiddleware.php
@@ -11,9 +11,8 @@
 
 namespace Symfony\Bridge\Doctrine\Messenger;
 
-use Doctrine\DBAL\Connection;
-use Doctrine\DBAL\Exception as DBALException;
 use Doctrine\ORM\EntityManagerInterface;
+use Symfony\Bridge\Doctrine\CheckConnection;
 use Symfony\Component\Messenger\Envelope;
 use Symfony\Component\Messenger\Middleware\StackInterface;
 use Symfony\Component\Messenger\Stamp\ConsumedByWorkerStamp;
@@ -38,24 +37,10 @@ private function pingConnection(EntityManagerInterface $entityManager): void
     {
         $connection = $entityManager->getConnection();
 
-        try {
-            $this->executeDummySql($connection);
-        } catch (DBALException) {
-            $connection->close();
-            // Attempt to reestablish the lazy connection by sending another query.
-            $this->executeDummySql($connection);
-        }
+        CheckConnection::checkConnection($connection);
 
         if (!$entityManager->isOpen()) {
             $this->managerRegistry->resetManager($this->entityManagerName);
         }
     }
-
-    /**
-     * @throws DBALException
-     */
-    private function executeDummySql(Connection $connection): void
-    {
-        $connection->executeQuery($connection->getDatabasePlatform()->getDummySelectSQL());
-    }
 }
diff --git a/src/Symfony/Bridge/Doctrine/Middleware/DoctrineConnectionDriver.php b/src/Symfony/Bridge/Doctrine/Middleware/DoctrineConnectionDriver.php
@@ -0,0 +1,70 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <fabien@symfony.com>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\Bridge\Doctrine\Middleware;
+
+use Doctrine\DBAL\Connection;
+use Doctrine\DBAL\Driver;
+use Doctrine\DBAL\Driver\Connection as DriverConnection;
+use Doctrine\DBAL\Driver\Middleware\AbstractDriverMiddleware;
+use Doctrine\ORM\EntityManagerInterface;
+use Doctrine\Persistence\ManagerRegistry;
+use Symfony\Bridge\Doctrine\CheckConnection;
+use Symfony\Component\DependencyInjection\ContainerInterface;
+
+class DoctrineConnectionDriver extends AbstractDriverMiddleware
+{
+    public function __construct(private Driver $driver, private readonly ManagerRegistry $managerRegistry, private readonly ContainerInterface $container)
+    {
+        parent::__construct($driver);
+    }
+
+    public function connect(array $params): DriverConnection
+    {
+        $connectionServices = $this->managerRegistry->getConnectionNames();
+
+        foreach ($connectionServices as $connectionServiceName) {
+            if (!$this->container->initialized($connectionServiceName)) {
+                continue;
+            }
+
+            $connection = $this->container->get($connectionServiceName);
+
+            if (!$connection instanceof Connection) {
+                continue;
+            }
+
+            if ($connection->isConnected()) {
+                CheckConnection::checkConnection($connection);
+            }
+
+            $managerNames = $this->managerRegistry->getManagerNames();
+
+            foreach ($managerNames as $managerName) {
+                if (!$this->container->initialized($managerName)) {
+                    continue;
+                }
+
+                $manager = $this->container->get($managerName);
+
+                if (!$manager instanceof EntityManagerInterface) {
+                    continue;
+                }
+
+                if (!$manager->isOpen()) {
+                    $this->managerRegistry->resetManager($managerName);
+                }
+            }
+        }
+
+        return parent::connect($params);
+    }
+}
diff --git a/src/Symfony/Bridge/Doctrine/Middleware/DoctrineConnectionMiddleware.php b/src/Symfony/Bridge/Doctrine/Middleware/DoctrineConnectionMiddleware.php
@@ -0,0 +1,32 @@
+<?php
+
+/*
+ * This file is part of the Symfony package.
+ *
+ * (c) Fabien Potencier <fabien@symfony.com>
+ *
+ * For the full copyright and license information, please view the LICENSE
+ * file that was distributed with this source code.
+ */
+
+namespace Symfony\Bridge\Doctrine\Middleware;
+
+use Doctrine\DBAL\Driver;
+use Doctrine\DBAL\Driver\Middleware;
+use Doctrine\Persistence\ManagerRegistry;
+use Symfony\Component\DependencyInjection\ContainerInterface;
+
+/**
+ * Based on https://github.com/Baldinof/roadrunner-bundle/blob/3.x/src/Integration/Doctrine/DoctrineORMMiddleware.php.
+ */
+class DoctrineConnectionMiddleware implements Middleware
+{
+    public function __construct(private ManagerRegistry $managerRegistry, private ContainerInterface $container)
+    {
+    }
+
+    public function wrap(Driver $driver): Driver
+    {
+        return new DoctrineConnectionDriver($driver, $this->managerRegistry, $this->container);
+    }
+}
