diff --git a/src/Database/PDO.php b/src/Database/PDO.php index 069ef88f8..245b0dfad 100644 --- a/src/Database/PDO.php +++ b/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; + } } /** diff --git a/tests/unit/PDOTest.php b/tests/unit/PDOTest.php index a17de1eb3..45e9a12a2 100644 --- a/tests/unit/PDOTest.php +++ b/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 {