Update WPCLI update/delete query for affected rows

Author: karthick-murugan

src/DB_Command.php

@@ -495,24 +495,38 @@ public function cli( $_, $assoc_args ) {
 	 *     +---------+-----------------------+
 	 */
 	public function query( $args, $assoc_args ) {
-
 		$command = sprintf( '/usr/bin/env mysql%s --no-auto-rehash', $this->get_defaults_flag_string( $assoc_args ) );
-		WP_CLI::debug( "Running shell command: {$command}", 'db' );
-
 		$assoc_args['database'] = DB_NAME;
-
-		// The query might come from STDIN.
+	
 		if ( ! empty( $args ) ) {
 			$assoc_args['execute'] = $args[0];
 		}
-
+	
 		if ( isset( $assoc_args['execute'] ) ) {
-			// Ensure that the SQL mode is compatible with WPDB.
 			$assoc_args['execute'] = $this->get_sql_mode_query( $assoc_args ) . $assoc_args['execute'];
 		}
-
-		WP_CLI::debug( 'Associative arguments: ' . json_encode( $assoc_args ), 'db' );
-		self::run( $command, $assoc_args );
+	
+		// Check if the query is an UPDATE or DELETE.
+		if ( isset( $assoc_args['execute'] ) && preg_match( '/\b(UPDATE|DELETE)\b/i', $assoc_args['execute'] ) ) {
+			// Append `SELECT ROW_COUNT()` to the query.
+			$assoc_args['execute'] .= '; SELECT ROW_COUNT();';
+		}
+	
+		list( $stdout, $stderr, $exit_code ) = self::run( $command, $assoc_args, false );
+	
+		if ( $exit_code ) {
+			WP_CLI::error( "Query failed: {$stderr}" );
+		}
+	
+		// For UPDATE/DELETE queries, parse the output to get the number of rows affected.
+		if ( isset( $assoc_args['execute'] ) && preg_match( '/\b(UPDATE|DELETE)\b/i', $assoc_args['execute'] ) ) {
+			$output_lines = explode( "\n", trim( $stdout ) );
+			$affected_rows = (int) trim( end( $output_lines ) );
+			WP_CLI::success( "Query succeeded. Rows affected: {$affected_rows}" );
+		} else {
+			WP_CLI::log( $stdout );
+			WP_CLI::success( 'Query executed successfully.' );
+		}
 	}
 
 	/**