From 2c04ace02236fb01c2aed9229d5da1ecab63e813 Mon Sep 17 00:00:00 2001
From: Jake Barnby <jakeb994@gmail.com>
Date: Wed, 8 Jan 2025 18:38:24 +1300
Subject: [PATCH] Add connection for project

---
 app/init.php   | 31 +++++++++++++++++++++++--------
 app/worker.php | 39 +++++++++++++++++++++------------------
 2 files changed, 44 insertions(+), 26 deletions(-)

diff --git a/app/init.php b/app/init.php
index b4e60d7924..51d9c38239 100644
--- a/app/init.php
+++ b/app/init.php
@@ -74,6 +74,7 @@ use Utopia\Logger\Adapter\Raygun;
 use Utopia\Logger\Adapter\Sentry;
 use Utopia\Logger\Log;
 use Utopia\Logger\Logger;
+use Utopia\Pools\Connection as PoolConnection;
 use Utopia\Pools\Group;
 use Utopia\Pools\Pool;
 use Utopia\Queue;
@@ -1412,7 +1413,26 @@ App::setResource('console', function () {
     ]);
 }, []);
 
-App::setResource('dbForProject', function (Group $pools, Database $dbForPlatform, Cache $cache, Document $project) {
+App::setResource('connectionForProject', function (Group $pools, Document $project) {
+    if ($project->isEmpty() || $project->getId() === 'console') {
+        return $pools
+            ->get('console')
+            ->pop();
+    }
+
+    try {
+        $dsn = new DSN($project->getAttribute('database'));
+    } catch (\InvalidArgumentException) {
+        // TODO: Temporary until all projects are using shared tables
+        $dsn = new DSN('mysql://' . $project->getAttribute('database'));
+    }
+
+    return $pools
+        ->get($dsn->getHost())
+        ->pop();
+}, ['pools', 'project']);
+
+App::setResource('dbForProject', function (Group $pools, PoolConnection $connectionForProject, Database $dbForPlatform, Cache $cache, Document $project) {
     if ($project->isEmpty() || $project->getId() === 'console') {
         return $dbForPlatform;
     }
@@ -1424,12 +1444,7 @@ App::setResource('dbForProject', function (Group $pools, Database $dbForPlatform
         $dsn = new DSN('mysql://' . $project->getAttribute('database'));
     }
 
-    $dbAdapter = $pools
-        ->get($dsn->getHost())
-        ->pop()
-        ->getResource();
-
-    $database = new Database($dbAdapter, $cache);
+    $database = new Database($connectionForProject->getResource(), $cache);
 
     $database
         ->setMetadata('host', \gethostname())
@@ -1452,7 +1467,7 @@ App::setResource('dbForProject', function (Group $pools, Database $dbForPlatform
     }
 
     return $database;
-}, ['pools', 'dbForPlatform', 'cache', 'project']);
+}, ['pools', 'connectionForProject', 'dbForPlatform', 'cache', 'project']);
 
 App::setResource('dbForPlatform', function (Group $pools, Cache $cache) {
     $dbAdapter = $pools
diff --git a/app/worker.php b/app/worker.php
index 6eb1363e9b..845af12efc 100644
--- a/app/worker.php
+++ b/app/worker.php
@@ -16,6 +16,7 @@ use Appwrite\Event\Migration;
 use Appwrite\Event\Usage;
 use Appwrite\Event\UsageDump;
 use Appwrite\Platform\Appwrite;
+use Swoole\Database\DetectsLostConnections;
 use Swoole\Runtime;
 use Utopia\Abuse\Adapters\TimeLimit\Redis as TimeLimitRedis;
 use Utopia\Cache\Adapter\Sharding;
@@ -25,11 +26,13 @@ use Utopia\Config\Config;
 use Utopia\Database\Database;
 use Utopia\Database\DateTime;
 use Utopia\Database\Document;
+use Utopia\Database\Exception as DatabaseException;
 use Utopia\Database\Validator\Authorization;
 use Utopia\DSN\DSN;
 use Utopia\Logger\Log;
 use Utopia\Logger\Logger;
 use Utopia\Platform\Service;
+use Utopia\Pools\Connection as PoolConnection;
 use Utopia\Pools\Group;
 use Utopia\Queue\Connection;
 use Utopia\Queue\Message;
@@ -66,13 +69,13 @@ Server::setResource('project', function (Message $message, Database $dbForPlatfo
     return $dbForPlatform->getDocument('projects', $project->getId());
 }, ['message', 'dbForPlatform']);
 
-Server::setResource('dbForProject', function (Cache $cache, Registry $register, Message $message, Document $project, Database $dbForPlatform) {
+Server::setResource('connectionForProject', function (Group $pools, Document $project) {
     if ($project->isEmpty() || $project->getId() === 'console') {
-        return $dbForPlatform;
+        return $pools
+            ->get('console')
+            ->pop();
     }
 
-    $pools = $register->get('pools');
-
     try {
         $dsn = new DSN($project->getAttribute('database'));
     } catch (\InvalidArgumentException) {
@@ -80,12 +83,17 @@ Server::setResource('dbForProject', function (Cache $cache, Registry $register,
         $dsn = new DSN('mysql://' . $project->getAttribute('database'));
     }
 
-    $adapter = $pools
+    return $pools
         ->get($dsn->getHost())
-        ->pop()
-        ->getResource();
+        ->pop();
+}, ['pools', 'project']);
 
-    $database = new Database($adapter, $cache);
+Server::setResource('dbForProject', function (PoolConnection $connectionForProject, Cache $cache, Registry $register, Message $message, Document $project, Database $dbForPlatform) {
+    if ($project->isEmpty() || $project->getId() === 'console') {
+        return $dbForPlatform;
+    }
+
+    $database = new Database($connectionForProject->getResource(), $cache);
 
     try {
         $dsn = new DSN($project->getAttribute('database'));
@@ -109,12 +117,12 @@ Server::setResource('dbForProject', function (Cache $cache, Registry $register,
     }
 
     return $database;
-}, ['cache', 'register', 'message', 'project', 'dbForPlatform']);
+}, ['connectionForProject', 'cache', 'register', 'message', 'project', 'dbForPlatform']);
 
-Server::setResource('getProjectDB', function (Group $pools, Database $dbForPlatform, $cache) {
+Server::setResource('getProjectDB', function (Group $pools, PoolConnection $connectionForProject, Database $dbForPlatform, $cache) {
     $databases = []; // TODO: @Meldiron This should probably be responsibility of utopia-php/pools
 
-    return function (Document $project) use ($pools, $dbForPlatform, $cache, &$databases): Database {
+    return function (Document $project) use ($pools, $connectionForProject, $dbForPlatform, $cache, &$databases): Database {
         if ($project->isEmpty() || $project->getId() === 'console') {
             return $dbForPlatform;
         }
@@ -146,12 +154,7 @@ Server::setResource('getProjectDB', function (Group $pools, Database $dbForPlatf
             return $database;
         }
 
-        $dbAdapter = $pools
-            ->get($dsn->getHost())
-            ->pop()
-            ->getResource();
-
-        $database = new Database($dbAdapter, $cache);
+        $database = new Database($connectionForProject->getResource(), $cache);
 
         $databases[$dsn->getHost()] = $database;
 
@@ -171,7 +174,7 @@ Server::setResource('getProjectDB', function (Group $pools, Database $dbForPlatf
 
         return $database;
     };
-}, ['pools', 'dbForPlatform', 'cache']);
+}, ['pools', 'connectionForProject', 'dbForPlatform', 'cache']);
 
 Server::setResource('abuseRetention', function () {
     return time() - (int) System::getEnv('_APP_MAINTENANCE_RETENTION_ABUSE', 86400);