Merge pull request #277 from karthick-murugan/wpcli/delet-update-changes

mrsdizzie · web-flow · commit 19b2555cca1b · 2025-03-10T11:01:50.000-04:00
Add Affected Rows Count for queries that can modify data in wp db query
diff --git a/features/db.feature b/features/db.feature
@@ -326,3 +326,23 @@ Feature: Perform database operations
       """
       latin1_spanish_ci
       """
+
+  Scenario: Row modifying queries should return the number of affected rows
+    Given a WP install
+    When I run `wp db query "UPDATE wp_users SET user_status = 1 WHERE ID = 1"`
+    Then STDOUT should contain:
+      """
+      Query succeeded. Rows affected: 1
+      """
+    
+    When I run `wp db query "SELECT * FROM wp_users WHERE ID = 1"`
+    Then STDOUT should not contain:
+      """
+      Rows affected
+      """
+
+    When I run `wp db query "DELETE FROM wp_users WHERE ID = 1"`
+    Then STDOUT should contain:
+      """
+      Query succeeded. Rows affected: 1
+      """
diff --git a/src/DB_Command.php b/src/DB_Command.php
@@ -511,8 +511,26 @@ public function query( $args, $assoc_args ) {
 			$assoc_args['execute'] = $this->get_sql_mode_query( $assoc_args ) . $assoc_args['execute'];
 		}
 
+		$is_row_modifying_query = isset( $assoc_args['execute'] ) && preg_match( '/\b(UPDATE|DELETE|INSERT|REPLACE|LOAD DATA)\b/i', $assoc_args['execute'] );
+
+		if ( $is_row_modifying_query ) {
+			$assoc_args['execute'] .= '; SELECT ROW_COUNT();';
+		}
+
 		WP_CLI::debug( 'Associative arguments: ' . json_encode( $assoc_args ), 'db' );
-		self::run( $command, $assoc_args );
+		list( $stdout, $stderr, $exit_code ) = self::run( $command, $assoc_args, false );
+
+		if ( $exit_code ) {
+			WP_CLI::error( "Query failed: {$stderr}" );
+		}
+
+		if ( $is_row_modifying_query ) {
+			$output_lines  = explode( "\n", trim( $stdout ) );
+			$affected_rows = (int) trim( end( $output_lines ) );
+			WP_CLI::success( "Query succeeded. Rows affected: {$affected_rows}" );
+		} elseif ( ! empty( $stdout ) ) {
+			WP_CLI::line( $stdout );
+		}
 	}
 
 	/**
