Merge pull request #124 from pmgarman/issue/108

Adds `--human-readable` parameter to `db size` command
wp-cli · Dec 17, 2018 · 2ab6bd0 · 2ab6bd0
2 parents f5662ed + 0af95c6
commit 2ab6bd0
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 3 deletions.
diff --git a/features/db-size.feature b/features/db-size.feature
@@ -30,6 +30,55 @@ Feature: Display database size
       wp_cli_test
       """
 
+  Scenario: Display only database size in a human readable format for a WordPress install
+    Given a WP install
+
+    When I run `wp db size --human-readable`
+    Then STDOUT should contain:
+      """
+      wp_cli_test
+      """
+
+    And STDOUT should contain:
+      """
+      KB
+      """
+
+    When I try `wp db size --human-readable --size_format=b`
+    Then the return code should not be 0
+    And STDERR should contain:
+      """
+      Cannot use --size_format and --human-readable arguments at the same time.
+      """
+    And STDOUT should be empty
+
+  Scenario: Display only table sizes in a human readable format for a WordPress install
+    Given a WP install
+
+    When I run `wp db size --tables --human-readable`
+    Then STDOUT should contain:
+      """
+      wp_posts
+      """
+
+    And STDOUT should contain:
+      """
+      KB
+      """
+
+    But STDOUT should not contain:
+      """
+      wp_cli_test
+      """
+
+    When I try `wp db size --tables --human-readable --size_format=b`
+    Then the return code should not be 0
+    And STDERR should contain:
+      """
+      Cannot use --size_format and --human-readable arguments at the same time.
+      """
+    And STDOUT should be empty
+
   Scenario: Display only database size in bytes for a WordPress install
     Given a WP install
 

diff --git a/src/DB_Command.php b/src/DB_Command.php
@@ -685,6 +685,9 @@ public function tables( $args, $assoc_args ) {
 	 * [--tables]
 	 * : Display each table name and size instead of the database size.
 	 *
+	 * [--human-readable]
+	 * : Display database sizes in human readable formats.
+	 *
 	 * [--format]
 	 * : table, csv, json
 	 * ---
@@ -751,11 +754,17 @@ public function size( $args, $assoc_args ) {
 
 		$format = WP_CLI\Utils\get_flag_value( $assoc_args, 'format' );
 		$size_format = WP_CLI\Utils\get_flag_value( $assoc_args, 'size_format' );
+		$human_readable = WP_CLI\Utils\get_flag_value( $assoc_args, 'human-readable', false );
 		$tables = WP_CLI\Utils\get_flag_value( $assoc_args, 'tables' );
 		$tables = ! empty( $tables );
 
+		if( ! is_null( $size_format ) && $human_readable ) {
+			WP_CLI::error( "Cannot use --size_format and --human-readable arguments at the same time." );
+		}
+
 		unset( $assoc_args['format'] );
 		unset( $assoc_args['size_format'] );
+		unset( $assoc_args['human-readable'] );
 		unset( $assoc_args['tables'] );
 
 		if ( empty( $args ) && empty( $assoc_args ) ) {
@@ -766,7 +775,7 @@ public function size( $args, $assoc_args ) {
 		$rows = array();
 		$fields = array( 'Name', 'Size' );
 
-		$default_unit = ( empty( $size_format ) ) ? ' B' : '';
+		$default_unit = ( empty( $size_format ) && ! $human_readable ) ? ' B' : '';
 
 		if ( $tables ) {
 
@@ -803,7 +812,7 @@ public function size( $args, $assoc_args ) {
 				);
 		}
 
-		if ( ! empty( $size_format ) ) {
+		if ( ! empty( $size_format ) || $human_readable ) {
 			foreach( $rows as $index => $row ) {
 					// These added WP 4.4.0.
 					if ( ! defined( 'KB_IN_BYTES' ) ) {
@@ -819,6 +828,13 @@ public function size( $args, $assoc_args ) {
 						define( 'TB_IN_BYTES', 1024 * GB_IN_BYTES );
 					}
 
+					if ( $human_readable ) {
+						$size_key = floor( log( $row['Size'] ) / log( 1000 ) );
+					    $sizes = array( 'B', 'KB', 'MB', 'GB', 'TB' );
+
+						$size_format = isset( $sizes[ $size_key ] ) ? $sizes[ $size_key ] : $sizes[ 0 ];
+					}
+
 					// Display the database size as a number.
 					switch( $size_format ) {
 						case 'TB':
@@ -869,7 +885,7 @@ public function size( $args, $assoc_args ) {
 			}
 		}
 
-		if ( ! empty( $size_format) && ! $tables && ! $format ) {
+		if ( ! empty( $size_format) && ! $tables && ! $format && ! $human_readable ) {
 			WP_CLI::Line( filter_var( $rows[0]['Size'], FILTER_SANITIZE_NUMBER_INT ) );
 		} else {
 			// Display the rows.