1 files changed, 89 insertions, 0 deletions
diff --git a/Echo/maintenance/recomputeNotifCounts.php b/Echo/maintenance/recomputeNotifCounts.php
new file mode 100644
index 00000000..23560772
--- /dev/null
+++ b/Echo/maintenance/recomputeNotifCounts.php
@@ -0,0 +1,89 @@
+<?php
+/**
+ * Recompute notification counts for all users.
+ *
+ * @ingroup Maintenance
+ */
+require_once getenv( 'MW_INSTALL_PATH' ) !== false
+	? getenv( 'MW_INSTALL_PATH' ) . '/maintenance/Maintenance.php'
+	: __DIR__ . '/../../../maintenance/Maintenance.php';
+
+/**
+ * Maintenance script that recomputes notification counts
+ *
+ * @ingroup Maintenance
+ */
+class RecomputeNotifCounts extends Maintenance {
+
+	public function __construct() {
+		parent::__construct();
+
+		$this->addDescription( 'Recompute notification counts for all users.' );
+		$this->addOption( 'user-ids',
+			'Comma-separated list of users to recompute notification counts for', false, true );
+		$this->addOption( 'notif-types',
+			'Recompute counts for all users who have unread notifications of one of these types (comma-separated)',
+			false, true );
+
+		$this->setBatchSize( 500 );
+
+		$this->requireExtension( 'Echo' );
+	}
+
+	public function execute() {
+		$dbFactory = MWEchoDbFactory::newFromDefault();
+		$dbwEcho = $dbFactory->getEchoDb( DB_MASTER );
+		$dbrEcho = $dbFactory->getEchoDb( DB_REPLICA );
+		$dbr = wfGetDB( DB_REPLICA );
+
+		$userIDs = $this->getOption( 'user-ids' );
+		$userIDs = $userIDs ? explode( ',', $userIDs ) : null;
+		$notifTypes = $this->getOption( 'notif-types' );
+		$notifTypes = $notifTypes ? explode( ',', $notifTypes ) : null;
+
+		if ( $userIDs ) {
+			$userIterator = array_chunk( $userIDs, $this->getBatchSize() );
+		} elseif ( $notifTypes ) {
+			$userIterator = new BatchRowIterator(
+				$dbrEcho,
+				[ 'echo_event', 'echo_notification' ],
+				'notification_user',
+				$this->getBatchSize()
+			);
+			$userIterator->addJoinConditions( [
+				'echo_notification' => [ 'INNER JOIN', 'notification_event=event_id' ]
+			] );
+			$userIterator->addConditions( [
+				'event_type' => $notifTypes,
+				'notification_read_timestamp' => null
+			] );
+			$userIterator->addOptions( [
+				'GROUP BY' => 'notification_user'
+			] );
+		} else {
+			$userQuery = User::getQueryInfo();
+			$userIterator = new BatchRowIterator( $dbr, $userQuery['tables'], 'user_id', $this->getBatchSize() );
+			$userIterator->setFetchColumns( $userQuery['fields'] );
+			$userIterator->addJoinConditions( $userQuery['joins'] );
+		}
+
+		$count = 0;
+		foreach ( $userIterator as $batch ) {
+			foreach ( $batch as $rowOrID ) {
+				if ( is_object( $rowOrID ) && isset( $rowOrID->user_id ) ) {
+					$user = User::newFromRow( $rowOrID );
+				} else {
+					$user = User::newFromId( is_object( $rowOrID ) ? $rowOrID->notification_user : $rowOrID );
+				}
+				$notifUser = MWEchoNotifUser::newFromUser( $user );
+				$notifUser->resetNotificationCount();
+			}
+			$count += count( $batch );
+			$this->output( "$count users' counts recomputed.\n" );
+			$dbFactory->waitForReplicas();
+		}
+	}
+}
+
+$maintClass = RecomputeNotifCounts::class;
+require_once RUN_MAINTENANCE_IF_MAIN;