fix: Added new db size method

vanguardbackup · Jul 27, 2024 · 308d378 · 308d378
1 parent b580e1c
commit 308d378
Show file tree

Hide file tree

Showing 3 changed files with 73 additions and 2 deletions.
diff --git a/app/Services/Backup/Backup.php b/app/Services/Backup/Backup.php
@@ -172,6 +172,54 @@ public function getRemoteDirectorySize(SFTPInterface $sftp, string $path): int
         return is_string($output) ? (int) trim($output) : 0;
     }
 
+    /**
+     * @throws SFTPConnectionException
+     */
+    public function getRemoteDatabaseSize(SFTPInterface $sftp, string $databaseType, string $databaseName, string $password): int
+    {
+        $this->logInfo('Getting remote database size.', ['database_type' => $databaseType, 'database' => $databaseName]);
+
+        $this->validateSFTP($sftp);
+
+        if ($databaseType === BackupConstants::DATABASE_TYPE_MYSQL) {
+            $sizeCommand = sprintf(
+                'mysql -p%s -e "SELECT ROUND(SUM(data_length + index_length) / 1024 / 1024, 1) AS size_mb FROM information_schema.tables WHERE table_schema = \'%s\';"',
+                escapeshellarg($password),
+                escapeshellarg($databaseName)
+            );
+
+        } elseif ($databaseType === BackupConstants::DATABASE_TYPE_POSTGRESQL) {
+            $sizeCommand = sprintf(
+                'PGPASSWORD=%s psql -d %s -c "SELECT pg_size_pretty(pg_database_size(\'%s\')) AS size;" -t | awk \'{print $1}\'',
+                escapeshellarg($password),
+                escapeshellarg($databaseName),
+                escapeshellarg($databaseName)
+            );
+
+        } else {
+            $this->logError('Unsupported database type.', ['database_type' => $databaseType]);
+            throw new SFTPConnectionException('Unsupported database type.');
+        }
+
+        $output = $sftp->exec($sizeCommand);
+        $this->logDebug('Database size command output.', ['command' => $sizeCommand, 'output' => $output]);
+
+        if (is_string($output)) {
+            $size = trim($output);
+            if ($databaseType === BackupConstants::DATABASE_TYPE_MYSQL) {
+                $sizeInMb = (float) $size;
+            } elseif ($databaseType === BackupConstants::DATABASE_TYPE_POSTGRESQL) {
+                $sizeInMb = $this->convertPgSizeToMb($size);
+            }
+
+            return (int) $sizeInMb;
+        }
+
+        $this->logError('Failed to get the database size.');
+
+        return 0;
+    }
+
     /**
      * @throws SFTPConnectionException
      */
@@ -585,4 +633,27 @@ protected function handleException(Exception $exception, string $context): void
     {
         $this->logError($context . ': ' . $exception->getMessage(), ['exception' => $exception]);
     }
+
+    private function convertPgSizeToMb(string $size): float
+    {
+        $units = [
+            'mb' => 1,
+            'kb' => 1 / 1024,
+            'gb' => 1024,
+            'tb' => 1024 * 1024,
+        ];
+
+        $size = strtolower($size);
+
+        preg_match('/([\d.]+)\s*([a-z]+)/', $size, $matches);
+
+        if (isset($matches[1]) && isset($matches[2])) {
+            $number = (float) $matches[1];
+            $unit = $matches[2];
+
+            return isset($units[$unit]) ? $number * $units[$unit] : $number / (1024 * 1024); // Assume bytes if unit is not recognized
+        }
+
+        return (float) $size / (1024 * 1024);
+    }
 }
diff --git a/app/Services/Backup/Tasks/DatabaseBackupTask.php b/app/Services/Backup/Tasks/DatabaseBackupTask.php
@@ -38,7 +38,7 @@ protected function performBackup(): void
         $dumpFileName = $this->generateBackupFileName('sql');
         $remoteDumpPath = '/tmp/' . $dumpFileName;
 
-        $dirSize = $this->getRemoteDirectorySize($sftp, $remoteDumpPath);
+        $dirSize = $this->getRemoteDatabaseSize($sftp, $databaseType, $databaseName, $databasePassword);
         $this->backupSize = $dirSize;
 
         $this->dumpRemoteDatabase($sftp, $databaseType, $remoteDumpPath, $databasePassword, $databaseName, $this->backupTask->getAttribute('excluded_database_tables'));

diff --git a/tests/Unit/Services/Backup/Tasks/DatabaseBackupTaskTest.php b/tests/Unit/Services/Backup/Tasks/DatabaseBackupTaskTest.php
@@ -146,7 +146,7 @@
     $this->databaseBackupTask->shouldReceive('dumpRemoteDatabase')->andReturnNull();
     $this->databaseBackupTask->shouldReceive('createBackupDestinationInstance')->andReturn($this->s3Mock);
     $this->databaseBackupTask->shouldReceive('rotateOldBackups')->once()->andReturnNull();
-    $this->databaseBackupTask->shouldReceive('getRemoteDirectorySize')->andReturn('500');
+    $this->databaseBackupTask->shouldReceive('getRemoteDatabaseSize')->andReturn('500');
 
     $this->databaseBackupTask->shouldReceive('updateBackupTaskStatus')->twice();