Add tests, move hooks to API Layer

app/controllers/api/databases.php

@@ -5,6 +5,7 @@ use Appwrite\Detector\Detector;
 use Appwrite\Event\Database as EventDatabase;
 use Appwrite\Event\Delete;
 use Appwrite\Event\Event;
+use Appwrite\Event\Usage;
 use Appwrite\Extend\Exception;
 use Appwrite\Network\Validator\Email;
 use Appwrite\Utopia\Database\Validator\CustomId;
@@ -452,7 +453,8 @@ App::post('/v1/databases')
     ->inject('response')
     ->inject('dbForProject')
     ->inject('queueForEvents')
-    ->action(function (string $databaseId, string $name, bool $enabled, Response $response, Database $dbForProject, Event $queueForEvents) {
+    ->inject('queueForUsage')
+    ->action(function (string $databaseId, string $name, bool $enabled, Response $response, Database $dbForProject, Event $queueForEvents, Usage $queueForUsage) {
 
         $databaseId = $databaseId == 'unique()' ? ID::unique() : $databaseId;
 
@@ -502,6 +504,7 @@ App::post('/v1/databases')
         }
 
         $queueForEvents->setParam('databaseId', $database->getId());
+        $queueForUsage->addMetric(str_replace(['{databaseInternalId}'], [$database->getInternalId()], METRIC_DATABASE_ID_STORAGE), 1); // per database
 
         $response
             ->setStatusCode(Response::STATUS_CODE_CREATED)
@@ -733,7 +736,8 @@ App::delete('/v1/databases/:databaseId')
     ->inject('dbForProject')
     ->inject('queueForDatabase')
     ->inject('queueForEvents')
-    ->action(function (string $databaseId, Response $response, Database $dbForProject, EventDatabase $queueForDatabase, Event $queueForEvents) {
+    ->inject('queueForUsage')
+    ->action(function (string $databaseId, Response $response, Database $dbForProject, EventDatabase $queueForDatabase, Event $queueForEvents, Usage $queueForUsage) {
 
         $database = $dbForProject->getDocument('databases', $databaseId);
 
@@ -756,6 +760,9 @@ App::delete('/v1/databases/:databaseId')
             ->setParam('databaseId', $database->getId())
             ->setPayload($response->output($database, Response::MODEL_DATABASE));
 
+        $queueForUsage
+            ->addMetric(METRIC_DATABASES_STORAGE, 1); // Global, deletion forces full recalculation
+
         $response->noContent();
     });
 
@@ -2350,7 +2357,8 @@ App::delete('/v1/databases/:databaseId/collections/:collectionId/attributes/:key
     ->inject('dbForProject')
     ->inject('queueForDatabase')
     ->inject('queueForEvents')
-    ->action(function (string $databaseId, string $collectionId, string $key, Response $response, Database $dbForProject, EventDatabase $queueForDatabase, Event $queueForEvents) {
+    ->inject('queueForUsage')
+    ->action(function (string $databaseId, string $collectionId, string $key, Response $response, Database $dbForProject, EventDatabase $queueForDatabase, Event $queueForEvents, Usage $queueForUsage) {
 
         $db = Authorization::skip(fn () => $dbForProject->getDocument('databases', $databaseId));
 
@@ -2435,6 +2443,9 @@ App::delete('/v1/databases/:databaseId/collections/:collectionId/attributes/:key
             ->setContext('database', $db)
             ->setPayload($response->output($attribute, $model));
 
+        $queueForUsage
+            ->addMetric(str_replace(['{databaseInternalId}', '{collectionInternalId}'], [$db->getInternalId(), $collection->getInternalId()], METRIC_DATABASE_ID_COLLECTION_ID_STORAGE), 1); // per collection
+
         $response->noContent();
     });
 
@@ -2810,8 +2821,9 @@ App::post('/v1/databases/:databaseId/collections/:collectionId/documents')
     ->inject('dbForProject')
     ->inject('user')
     ->inject('queueForEvents')
+    ->inject('queueForUsage')
     ->inject('mode')
-    ->action(function (string $databaseId, string $documentId, string $collectionId, string|array $data, ?array $permissions, Response $response, Database $dbForProject, Document $user, Event $queueForEvents, string $mode) {
+    ->action(function (string $databaseId, string $documentId, string $collectionId, string|array $data, ?array $permissions, Response $response, Database $dbForProject, Document $user, Event $queueForEvents, Usage $queueForUsage, string $mode) {
 
         $data = (\is_string($data)) ? \json_decode($data, true) : $data; // Cast to JSON array
 
@@ -3027,6 +3039,9 @@ App::post('/v1/databases/:databaseId/collections/:collectionId/documents')
             ->setContext('database', $database)
             ->setPayload($response->getPayload(), sensitive: $relationships);
 
+        
+        $queueForUsage
+            ->addMetric(str_replace(['{databaseInternalId}', '{collectionInternalId}'], [$database->getInternalId(), $collection->getInternalId()], METRIC_DATABASE_ID_COLLECTION_ID_STORAGE), 1); // per collection
     });
 
 App::get('/v1/databases/:databaseId/collections/:collectionId/documents')
@@ -3643,8 +3658,9 @@ App::delete('/v1/databases/:databaseId/collections/:collectionId/documents/:docu
     ->inject('dbForProject')
     ->inject('queueForDeletes')
     ->inject('queueForEvents')
+    ->inject('queueForUsage')
     ->inject('mode')
-    ->action(function (string $databaseId, string $collectionId, string $documentId, ?\DateTime $requestTimestamp, Response $response, Database $dbForProject, Delete $queueForDeletes, Event $queueForEvents, string $mode) {
+    ->action(function (string $databaseId, string $collectionId, string $documentId, ?\DateTime $requestTimestamp, Response $response, Database $dbForProject, Delete $queueForDeletes, Event $queueForEvents, Usage $queueForUsage, string $mode) {
         $database = Authorization::skip(fn () => $dbForProject->getDocument('databases', $databaseId));
 
         $isAPIKey = Auth::isAppUser(Authorization::getRoles());
@@ -3729,6 +3745,9 @@ App::delete('/v1/databases/:databaseId/collections/:collectionId/documents/:docu
             ->setContext('database', $database)
             ->setPayload($response->output($document, Response::MODEL_DOCUMENT), sensitive: $relationships);
 
+        $queueForUsage
+            ->addMetric(str_replace(['{databaseInternalId}', '{collectionInternalId}'], [$database->getInternalId(), $collection->getInternalId()], METRIC_DATABASE_ID_COLLECTION_ID_STORAGE), 1); // per collection
+
         $response->noContent();
     });
 

app/controllers/api/project.php

@@ -225,23 +225,6 @@ App::get('/v1/project/usage')
             ];
         }, $dbForProject->find('functions'));
 
-        $databasesStorageBreakdown = array_map(function ($database) use ($dbForProject) {
-            $id = $database->getId();
-            $name = $database->getAttribute('name');
-            $metric = str_replace('{databaseInternalId}', $database->getInternalId(), METRIC_DATABASE_ID_STORAGE);
-
-            $value = $dbForProject->findOne('stats', [
-                Query::equal('metric', [$metric]),
-                Query::equal('period', ['inf'])
-            ]);
-
-            return [
-                'resourceId' => $id,
-                'name' => $name,
-                'value' => $value['value'] ?? 0,
-            ];
-        }, $dbForProject->find('databases'));
-
         $executionsMbSecondsBreakdown = array_map(function ($function) use ($dbForProject) {
             $id = $function->getId();
             $name = $function->getAttribute('name');

app/controllers/shared/api.php

@@ -83,7 +83,6 @@ $databaseListener = function (string $event, Document $document, Document $proje
             break;
         case $document->getCollection() === 'databases': // databases
             $queueForUsage
-                ->addMetric(METRIC_DATABASES_STORAGE, 1)
                 ->addMetric(METRIC_DATABASES, $value); // per project
 
             if ($event === Database::EVENT_DOCUMENT_DELETE) {
@@ -96,9 +95,7 @@ $databaseListener = function (string $event, Document $document, Document $proje
             $databaseInternalId = $parts[1] ?? 0;
             $queueForUsage
                 ->addMetric(METRIC_COLLECTIONS, $value) // per project
-                ->addMetric(METRIC_DATABASES_STORAGE, 1)
                 ->addMetric(str_replace('{databaseInternalId}', $databaseInternalId, METRIC_DATABASE_ID_COLLECTIONS), $value)
-                ->addMetric(str_replace('{databaseInternalId}', $databaseInternalId, METRIC_DATABASE_ID_STORAGE), 1); // per database
             ;
 
             if ($event === Database::EVENT_DOCUMENT_DELETE) {
@@ -113,8 +110,7 @@ $databaseListener = function (string $event, Document $document, Document $proje
             $queueForUsage
                 ->addMetric(METRIC_DOCUMENTS, $value)  // per project
                 ->addMetric(str_replace('{databaseInternalId}', $databaseInternalId, METRIC_DATABASE_ID_DOCUMENTS), $value) // per database
-                ->addMetric(str_replace(['{databaseInternalId}', '{collectionInternalId}'], [$databaseInternalId, $collectionInternalId], METRIC_DATABASE_ID_COLLECTION_ID_DOCUMENTS), $value)  // per collection
-                ->addMetric(str_replace(['{databaseInternalId}', '{collectionInternalId}'], [$databaseInternalId, $collectionInternalId], METRIC_DATABASE_ID_COLLECTION_ID_STORAGE), 1); // per collection
+                ->addMetric(str_replace(['{databaseInternalId}', '{collectionInternalId}'], [$databaseInternalId, $collectionInternalId], METRIC_DATABASE_ID_COLLECTION_ID_DOCUMENTS), $value);  // per collection
             break;
         case $document->getCollection() === 'buckets': //buckets
             $queueForUsage

src/Appwrite/Platform/Workers/UsageDump.php

@@ -71,7 +71,7 @@ class UsageDump extends Action
                         continue;
                     }
 
-                    if (str_ends_with($key, '.db_storage')) {
+                    if (str_ends_with($key, '.db_storage') && $value === 1) {
                         $this->handleDBStorageCalculation($key, $dbForProject);
                         return;
                     }
@@ -117,6 +117,9 @@ class UsageDump extends Action
     private function handleDBStorageCalculation(string $key, Database $dbForProject): void
     {
         $data = explode('.', $key);
+        $start = microtime(true);
+
+        var_dump('Calculating DB Storage for ' . $key);
 
         $updateMetric = function (Database $dbForProject, int $value, string $key, string $period, string|null $time) {
             $id = \md5("{$time}_{$period}_{$key}");
@@ -176,6 +179,8 @@ class UsageDump extends Action
                         break;
                     }
 
+                    var_dump('Calculated collection level, diff was ' . $diff . ' for ' . $key);
+
                     // Update Collection
                     $updateMetric($dbForProject, $diff, $key, $period, $time);
 
@@ -187,7 +192,7 @@ class UsageDump extends Action
                     $projectKey = 'db_storage';
                     $updateMetric($dbForProject, $diff, $projectKey, $period, $time);
                     break;
-                    // Database Level
+                // Database Level
                 case 2:
                     $databaseInternalId = $data[0];
                     $collections = $dbForProject->find('database_' . $databaseInternalId);
@@ -198,6 +203,12 @@ class UsageDump extends Action
 
                     $diff = $value - $previousValue;
 
+                    if ($diff === 0) {
+                        break;
+                    }
+
+                    var_dump('Calculated database level, diff was ' . $diff . ' for ' . $key);
+
                     // Update Database
                     $databaseKey = $data[0] . '.db_storage';
                     $updateMetric($dbForProject, $diff, $databaseKey, $period, $time);
@@ -222,11 +233,17 @@ class UsageDump extends Action
 
                     $diff = $value - $previousValue;
 
+                    var_dump('Calculated project level, diff was ' . $diff . ' for ' . $key);
+
                     // Update Project
                     $projectKey = 'db_storage';
                     $updateMetric($dbForProject, $diff, $projectKey, $period, $time);
                     break;
             }
         }
+
+        $end = microtime(true);
+
+        console::log('[' . DateTime::now() . '] DB Storage Calculation [' . $key . '] took ' . (($end - $start) * 1000) . ' milliseconds');
     }
 }

tests/e2e/General/UsageTest.php

@@ -590,6 +590,166 @@ class UsageTest extends Scope
         return $data;
     }
 
+    public function testDatabaseStoragePrepare(): array
+    {
+        $response = $this->client->call(
+            Client::METHOD_POST,
+            '/databases',
+            array_merge([
+                'content-type' => 'application/json',
+                'x-appwrite-project' => $this->getProject()['$id']
+            ], $this->getHeaders()),
+            [
+                'databaseId' => 'unique()',
+                'name' => 'dbStorageStats',
+            ]
+        );
+
+        $this->assertNotEmpty($response['body']['$id']);
+        $databaseId = $response['body']['$id'];
+
+        $response = $this->client->call(
+            Client::METHOD_POST,
+            '/databases/' . $databaseId . '/collections',
+            array_merge([
+                'content-type' => 'application/json',
+                'x-appwrite-project' => $this->getProject()['$id']
+            ], $this->getHeaders()),
+            [
+                'collectionId' => 'unique()',
+                'name' => 'collectionStorageStats',
+                'documentSecurity' => false,
+                'permissions' => [
+                    Permission::read(Role::any()),
+                    Permission::create(Role::any()),
+                    Permission::update(Role::any()),
+                    Permission::delete(Role::any()),
+                ],
+            ]
+        );
+
+        $this->assertNotEmpty($response['body']['$id']);
+        $collectionId = $response['body']['$id'];
+
+        $response = $this->client->call(
+            Client::METHOD_POST,
+            '/databases/' . $databaseId . '/collections/' . $collectionId . '/attributes' . '/string',
+            array_merge([
+                'content-type' => 'application/json',
+                'x-appwrite-project' => $this->getProject()['$id']
+            ], $this->getHeaders()),
+            [
+                'key' => 'data',
+                'size' => 100000,
+                'required' => true,
+            ]
+        );
+
+        return [
+            'databaseId' => $databaseId,
+            'collectionId' => $collectionId,
+        ];
+    }
+
+    /** @depends testDatabaseStoragePrepare */
+    #[Retry(count: 1)]
+    public function testDatabaseStorageStatsCreateDocument(array $data): array
+    {
+        $databaseId = $data['databaseId'];
+        $collectionId = $data['collectionId'];
+
+        $originalProjectMetrics = $this->client->call(
+            Client::METHOD_GET,
+            '/project/usage',
+            $this->getConsoleHeaders(),
+            [
+                'period' => '1d',
+                'startDate' => self::getToday(),
+                'endDate' => self::getTomorrow(),
+            ]
+        );
+
+        $this->assertEquals(200, $originalProjectMetrics['headers']['status-code']);
+        $this->assertArrayHasKey('databasesStorageTotal', $originalProjectMetrics['body']);
+
+        $originalProjectMetrics = $originalProjectMetrics['body'];
+
+        $originalDatabaseMetrics = $this->client->call(
+            Client::METHOD_GET,
+            '/databases/' . $databaseId . '/usage?range=30d',
+            $this->getConsoleHeaders()
+        );
+
+        $this->assertEquals(200, $originalDatabaseMetrics['headers']['status-code']);
+        $this->assertArrayHasKey('storageTotal', $originalDatabaseMetrics['body']);
+        $originalDatabaseMetrics = $originalDatabaseMetrics['body'];
+
+        // Create documents
+        for ($i = 0; $i < 100; $i++) {
+            $response = $this->client->call(
+                Client::METHOD_POST,
+                '/databases/' . $databaseId . '/collections/' . $collectionId . '/documents',
+                array_merge([
+                    'content-type' => 'application/json',
+                    'x-appwrite-project' => $this->getProject()['$id']
+                ], $this->getHeaders()),
+                [
+                    'documentId' => 'unique()',
+                    'data' => ['data' => str_repeat('a', 10000)],
+                ]
+            );
+
+            $this->assertEquals(201, $response['headers']['status-code']);
+        }
+
+        sleep(self::WAIT);
+
+        for ($i = 0; $i < 3; $i++) {
+            try {
+                $newProjectMetrics = $this->client->call(
+                    Client::METHOD_GET,
+                    '/project/usage',
+                    $this->getConsoleHeaders(),
+                    [
+                        'period' => '1d',
+                        'startDate' => self::getToday(),
+                        'endDate' => self::getTomorrow(),
+                    ]
+                );
+        
+                $this->assertEquals(200, $newProjectMetrics['headers']['status-code']);
+                $this->assertArrayHasKey('databasesStorageTotal', $newProjectMetrics['body']);
+                $this->assertGreaterThan($originalProjectMetrics['databasesStorageTotal'], $newProjectMetrics['body']['databasesStorageTotal']);
+        
+                $newProjectMetrics = $newProjectMetrics['body'];
+        
+                $newDatabaseMetrics = $this->client->call(
+                    Client::METHOD_GET,
+                    '/databases/' . $databaseId . '/usage?range=30d',
+                    $this->getConsoleHeaders()
+                );
+        
+                $this->assertEquals(200, $newDatabaseMetrics['headers']['status-code']);
+                $this->assertArrayHasKey('storageTotal', $newDatabaseMetrics['body']);
+                $this->assertGreaterThan($originalDatabaseMetrics['storageTotal'], $newDatabaseMetrics['body']['storageTotal']);
+
+                $newDatabaseMetrics = $newDatabaseMetrics['body'];
+
+                return [
+                    'databaseId' => $databaseId,
+                    'collectionId' => $collectionId,
+                    'currentProjectMetrics' => $newProjectMetrics,
+                    'currentDatabaseMetrics' => $newDatabaseMetrics,
+                ];
+            } catch (ExpectationFailedException $e) {
+                if ($i === 2) {
+                    throw $e;
+                }
+                sleep(self::WAIT);
+                continue;
+            }
+        }
+    }
 
     /** @depends testDatabaseStats */
     public function testPrepareFunctionsStats(array $data): array