Merge pull request #705 from utopia-php/feat-enable-reconnect

Enable reconnection + retry

Author: abnegate

src/Database/PDO.php

@@ -3,6 +3,7 @@
 namespace Utopia\Database;
 
 use InvalidArgumentException;
+use Utopia\CLI\Console;
 
 /**
  * A PDO wrapper that forwards method calls to the internal PDO instance.
@@ -41,7 +42,27 @@ public function __construct(
      */
     public function __call(string $method, array $args): mixed
     {
-        return $this->pdo->{$method}(...$args);
+        try {
+            return $this->pdo->{$method}(...$args);
+        } catch (\Throwable $e) {
+            if (Connection::hasError($e)) {
+                Console::warning('[Database] ' . $e->getMessage());
+                Console::warning('[Database] Lost connection detected. Reconnecting...');
+
+                $inTransaction = $this->pdo->inTransaction();
+
+                // Attempt to reconnect
+                $this->reconnect();
+
+                // If we weren't in a transaction, also retry the query
+                // In a transaction we can't retry as the state is attached to the previous connection
+                if (!$inTransaction) {
+                    return $this->pdo->{$method}(...$args);
+                }
+            }
+
+            throw $e;
+        }
     }
 
     /**

tests/unit/PDOTest.php

@@ -41,44 +41,44 @@ public function testMethodCallIsForwardedToPDO(): void
         $this->assertSame($pdoStatementMock, $result);
     }
 
-    //    public function testLostConnectionRetriesCall(): void
-    //    {
-    //        $dsn = 'sqlite::memory:';
-    //        $pdoWrapper = $this->getMockBuilder(PDO::class)
-    //            ->setConstructorArgs([$dsn, null, null, []])
-    //            ->onlyMethods(['reconnect'])
-    //            ->getMock();
-    //
-    //        $pdoMock = $this->getMockBuilder(\PDO::class)
-    //            ->disableOriginalConstructor()
-    //            ->getMock();
-    //        $pdoStatementMock = $this->getMockBuilder(\PDOStatement::class)
-    //            ->disableOriginalConstructor()
-    //            ->getMock();
-    //
-    //        $pdoMock->expects($this->exactly(2))
-    //            ->method('query')
-    //            ->with('SELECT 1')
-    //            ->will($this->onConsecutiveCalls(
-    //                $this->throwException(new \Exception("Lost connection")),
-    //                $pdoStatementMock
-    //            ));
-    //
-    //        $reflection = new ReflectionClass($pdoWrapper);
-    //        $pdoProperty = $reflection->getProperty('pdo');
-    //        $pdoProperty->setAccessible(true);
-    //        $pdoProperty->setValue($pdoWrapper, $pdoMock);
-    //
-    //        $pdoWrapper->expects($this->once())
-    //            ->method('reconnect')
-    //            ->willReturnCallback(function () use ($pdoWrapper, $pdoMock, $pdoProperty) {
-    //                $pdoProperty->setValue($pdoWrapper, $pdoMock);
-    //            });
-    //
-    //        $result = $pdoWrapper->query('SELECT 1');
-    //
-    //        $this->assertSame($pdoStatementMock, $result);
-    //    }
+    public function testLostConnectionRetriesCall(): void
+    {
+        $dsn = 'sqlite::memory:';
+        $pdoWrapper = $this->getMockBuilder(PDO::class)
+            ->setConstructorArgs([$dsn, null, null, []])
+            ->onlyMethods(['reconnect'])
+            ->getMock();
+
+        $pdoMock = $this->getMockBuilder(\PDO::class)
+            ->disableOriginalConstructor()
+            ->getMock();
+        $pdoStatementMock = $this->getMockBuilder(\PDOStatement::class)
+            ->disableOriginalConstructor()
+            ->getMock();
+
+        $pdoMock->expects($this->exactly(2))
+            ->method('query')
+            ->with('SELECT 1')
+            ->will($this->onConsecutiveCalls(
+                $this->throwException(new \Exception("Lost connection")),
+                $pdoStatementMock
+            ));
+
+        $reflection = new ReflectionClass($pdoWrapper);
+        $pdoProperty = $reflection->getProperty('pdo');
+        $pdoProperty->setAccessible(true);
+        $pdoProperty->setValue($pdoWrapper, $pdoMock);
+
+        $pdoWrapper->expects($this->once())
+            ->method('reconnect')
+            ->willReturnCallback(function () use ($pdoWrapper, $pdoMock, $pdoProperty) {
+                $pdoProperty->setValue($pdoWrapper, $pdoMock);
+            });
+
+        $result = $pdoWrapper->query('SELECT 1');
+
+        $this->assertSame($pdoStatementMock, $result);
+    }
 
     public function testNonLostConnectionExceptionIsRethrown(): void
     {