From f6607e88182de2742fe117998271dd59e723e364 Mon Sep 17 00:00:00 2001
From: Arlo Breault <abreault@wikimedia.org>
Date: Tue, 2 Aug 2022 14:34:16 -0400
Subject: [PATCH] Stop using wfGetDB

This is pulled out of I59be7cabb80ace98da3c7f6f36a0d3d4f6b17d23

Change-Id: Iab6a47320995e9adb1666cd0bb728f516a2fde69
---
 includes/Database.php     | 19 ++++++++++++++-----
 includes/TotalsLookup.php |  4 +++-
 2 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/includes/Database.php b/includes/Database.php
index 91210385..fccf9cf3 100644
--- a/includes/Database.php
+++ b/includes/Database.php
@@ -24,6 +24,7 @@ use FormatJson;
 use MediaWiki\Logger\LoggerFactory;
 use MediaWiki\MediaWikiServices;
 use WikiMap;
+use Wikimedia\Rdbms\DBConnRef;
 
 /**
  * Database logic
@@ -68,6 +69,14 @@ class Database {
 		$this->categoryManager = new CategoryManager();
 	}
 
+	/**
+	 * @param int $mode DB_PRIMARY or DB_REPLICA
+	 * @return DBConnRef
+	 */
+	public static function getDBConnectionRef( int $mode ): DBConnRef {
+		return MediaWikiServices::getInstance()->getDBLoadBalancer()->getConnectionRef( $mode );
+	}
+
 	/**
 	 * Get a specific LintError by id
 	 *
@@ -75,7 +84,7 @@ class Database {
 	 * @return bool|LintError
 	 */
 	public function getFromId( $id ) {
-		$row = wfGetDB( DB_REPLICA )->selectRow(
+		$row = self::getDBConnectionRef( DB_REPLICA )->selectRow(
 			'linter',
 			[ 'linter_cat', 'linter_params', 'linter_start', 'linter_end' ],
 			[ 'linter_id' => $id, 'linter_page' => $this->pageId ],
@@ -121,7 +130,7 @@ class Database {
 	 * @return LintError[]
 	 */
 	public function getForPage() {
-		$rows = wfGetDB( DB_REPLICA )->select(
+		$rows = self::getDBConnectionRef( DB_REPLICA )->select(
 			'linter',
 			[
 				'linter_id', 'linter_cat', 'linter_start',
@@ -214,7 +223,7 @@ class Database {
 	 */
 	public function setForPage( $errors ) {
 		$previous = $this->getForPage();
-		$dbw = wfGetDB( DB_PRIMARY );
+		$dbw = self::getDBConnectionRef( DB_PRIMARY );
 		if ( !$previous && !$errors ) {
 			return [ 'deleted' => [], 'added' => [] ];
 		} elseif ( !$previous && $errors ) {
@@ -290,7 +299,7 @@ class Database {
 	 * @return int
 	 */
 	private function getTotalsEstimate( $catId ) {
-		$dbr = wfGetDB( DB_REPLICA );
+		$dbr = self::getDBConnectionRef( DB_REPLICA );
 		// First see if there are no rows, or a moderate number
 		// within the limit specified by the MAX_ACCURATE_COUNT.
 		// The distinction between 0, a few and a lot is important
@@ -328,7 +337,7 @@ class Database {
 	 * @return int[]
 	 */
 	private function getTotalsAccurate( $conds = [] ) {
-		$rows = wfGetDB( DB_REPLICA )->select(
+		$rows = self::getDBConnectionRef( DB_REPLICA )->select(
 			'linter',
 			[ 'linter_cat', 'COUNT(*) AS count' ],
 			$conds,
diff --git a/includes/TotalsLookup.php b/includes/TotalsLookup.php
index 757c76ee..0cac672f 100644
--- a/includes/TotalsLookup.php
+++ b/includes/TotalsLookup.php
@@ -70,7 +70,9 @@ class TotalsLookup {
 				$this->makeKey( $cat ),
 				WANObjectCache::TTL_INDEFINITE,
 				static function ( $oldValue, &$ttl, &$setOpts, $oldAsOf ) use ( $cat, &$fetchedTotals ) {
-					$setOpts += MWDatabase::getCacheSetOptions( wfGetDB( DB_REPLICA ) );
+					$setOpts += MWDatabase::getCacheSetOptions(
+						Database::getDBConnectionRef( DB_REPLICA )
+					);
 					if ( $fetchedTotals === false ) {
 						$fetchedTotals = ( new Database( 0 ) )->getTotals();
 					}