From b1a89e2dd349710f63f3b128c22f480afbaf08ec Mon Sep 17 00:00:00 2001
From: Ephraim G <ephraim.gregor@10up.com>
Date: Fri, 5 Oct 2018 18:07:33 -0400
Subject: [PATCH] Fixes #27: Adds a progress bar

---
 src/classes/Log.php |  9 +++++
 src/classes/S3.php  | 84 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 93 insertions(+)

diff --git a/src/classes/Log.php b/src/classes/Log.php
index c342671..0461b61 100644
--- a/src/classes/Log.php
+++ b/src/classes/Log.php
@@ -128,6 +128,15 @@ public function getVerbosity() {
 		return $this->verbosity;
 	}
 
+	/**
+	 * Gets a reference to the output.
+	 *
+	 * @return OutputInterface
+	 */
+	public function getOutput() {
+		return $this->output;
+	}
+
 	/**
 	 * Return singleton instance of class
 	 *
diff --git a/src/classes/S3.php b/src/classes/S3.php
index 4e6d5b1..2e2d064 100644
--- a/src/classes/S3.php
+++ b/src/classes/S3.php
@@ -7,6 +7,7 @@
 
 namespace WPSnapshots;
 
+use Symfony\Component\Console\Helper\ProgressBar;
 use \Aws\S3\S3Client;
 use \Aws\Exception\AwsException;
 use WPSnapshots\Utils;
@@ -36,6 +37,13 @@ class S3 {
 	 */
 	private $region;
 
+	/**
+	 * Progress bar reference.
+	 *
+	 * @var \cli\progress\Bar
+	 */
+	private $progress_bar;
+
 	/**
 	 * Setup S3 client
 	 *
@@ -56,6 +64,23 @@ public function __construct( $config ) {
 
 		$this->repository = $config['repository'];
 		$this->region     = $config['region'];
+
+		ProgressBar::setPlaceholderFormatterDefinition(
+			'cur_bytes',
+			function ( $progressBar, $output ) {
+				return Utils\format_bytes( $progressBar->getProgress() );
+			}
+		);
+		ProgressBar::setPlaceholderFormatterDefinition(
+			'max_bytes',
+			function ( $progressBar, $output ) {
+				return Utils\format_bytes( $progressBar->getMaxSteps() );
+			}
+		);
+		ProgressBar::setFormatDefinition(
+			's3',
+			'%cur_bytes%/%max_bytes% [%bar%] %percent:3s%%'
+		);
 	}
 
 	/**
@@ -126,21 +151,31 @@ public function putSnapshot( $id, $project, $db_path, $files_path ) {
 	 */
 	public function downloadSnapshot( $id, $project, $db_path, $files_path ) {
 		try {
+			Log::instance()->write( 'Downloading database...' );
 			$db_download = $this->client->getObject(
 				[
 					'Bucket' => self::getBucketName( $this->repository ),
 					'Key'    => $project . '/' . $id . '/data.sql.gz',
 					'SaveAs' => $db_path,
+					'@http'  => [
+						'progress' => [ $this, 'progress' ],
+					],
 				]
 			);
+			$this->reset_progress();
 
+			Log::instance()->write( 'Downloading files...' );
 			$files_download = $this->client->getObject(
 				[
 					'Bucket' => self::getBucketName( $this->repository ),
 					'Key'    => $project . '/' . $id . '/files.tar.gz',
 					'SaveAs' => $files_path,
+					'@http'  => [
+						'progress' => [ $this, 'progress' ],
+					],
 				]
 			);
+			$this->reset_progress();
 		} catch ( \Exception $e ) {
 			$error = [
 				'message'        => $e->getMessage(),
@@ -294,4 +329,53 @@ public function createBucket() {
 
 		return true;
 	}
+
+	/**
+	 * Progress callback.
+	 *
+	 * @see https://docs.aws.amazon.com/sdk-for-php/v3/developer-guide/guide_configuration.html#progress
+	 *
+	 * @param int $expected_download_bytes Total expected bytes to be downloaded.
+	 * @param int $downloaded_bytes        Downloaded bytes so far.
+	 * @param int $expected_upload_bytes   Total expected bytes to be uploaded.
+	 * @param int $uploaded_bytes          Uploaded bytes so far.
+	 * @return void
+	 */
+	public function progress(
+		$expected_download_bytes,
+		$downloaded_bytes,
+		$expected_upload_bytes,
+		$uploaded_bytes
+	) {
+		if ( ! $this->progress_bar ) {
+			$this->progress_bar = new ProgressBar( Log::instance()->getOutput() );
+			$this->progress_bar->setFormat( 's3' );
+		}
+		if ( ! $this->progress_bar->getMaxSteps() ) {
+			if ( $expected_download_bytes ) {
+				$this->progress_bar->setMaxSteps( $expected_download_bytes );
+			} elseif ( $expected_upload_bytes ) {
+				$this->progress_bar->setMaxSteps( $expected_upload_bytes );
+			}
+		}
+
+		if ( $downloaded_bytes ) {
+			$this->progress_bar->setProgress( $downloaded_bytes );
+		} elseif ( $uploaded_bytes ) {
+			$this->progress_bar->setProgress( $uploaded_bytes );
+		}
+	}
+
+	/**
+	 * Resets the progress bar state.
+	 */
+	protected function reset_progress() {
+		if ( ! $this->progress_bar ) {
+			return;
+		}
+
+		$this->progress_bar->finish();
+		$this->progress_bar = null;
+		Log::instance()->write( '' );
+	}
 }