Make wfGetDB() return a MaintainableDBConnRef instance

This enforces the DB_* role checks of DBConnRef in more places

Depends-on: I9328e709fe5d81099338a31deef24d34db22d784
Change-Id: I0d7dacee3ec4ef67dc0b0f6551ad046c74dc47dc

includes/GlobalFunctions.php

@@ -2563,10 +2563,10 @@ function wfWikiID() {
  * @todo Replace calls to wfGetDB with calls to LoadBalancer::getConnection()
  *       on an injected instance of LoadBalancer.
  *
- * @return \Wikimedia\Rdbms\Database
+ * @return \Wikimedia\Rdbms\DBConnRef
  */
 function wfGetDB( $db, $groups = [], $wiki = false ) {
-	return wfGetLB( $wiki )->getConnection( $db, $groups, $wiki );
+	return wfGetLB( $wiki )->getMaintenanceConnectionRef( $db, $groups, $wiki );
 }
 
 /**

tests/phpunit/includes/Revision/RevisionStoreTest.php

@@ -12,6 +12,8 @@ use MediaWiki\Revision\RevisionStore;
 use MediaWiki\Revision\SlotRoleRegistry;
 use MediaWiki\Revision\SlotRecord;
 use MediaWiki\Storage\SqlBlobStore;
+use Wikimedia\Rdbms\ILoadBalancer;
+use Wikimedia\Rdbms\MaintainableDBConnRef;
 use MediaWikiTestCase;
 use MWException;
 use Title;
@@ -77,6 +79,17 @@ class RevisionStoreTest extends MediaWikiTestCase {
 			->disableOriginalConstructor()->getMock();
 	}
 
+	/**
+	 * @param ILoadBalancer $mockLoadBalancer
+	 * @param Database $db
+	 * @return callable
+	 */
+	private function getMockDBConnRefCallback( ILoadBalancer $mockLoadBalancer, IDatabase $db ) {
+		return function ( $i, $g, $domain, $flg ) use ( $mockLoadBalancer, $db ) {
+			return new MaintainableDBConnRef( $mockLoadBalancer, $db, $i );
+		};
+	}
+
 	/**
 	 * @return \PHPUnit_Framework_MockObject_MockObject|SqlBlobStore
 	 */
@@ -158,10 +171,14 @@ class RevisionStoreTest extends MediaWikiTestCase {
 		$this->setService( 'DBLoadBalancer', $mockLoadBalancer );
 
 		$db = $this->getMockDatabase();
-		// Title calls wfGetDB() which uses a regular Connection
+		// RevisionStore uses getConnectionRef
+		$mockLoadBalancer->expects( $this->any() )
+			->method( 'getConnectionRef' )
+			->willReturnCallback( $this->getMockDBConnRefCallback( $mockLoadBalancer, $db ) );
+		// Title calls wfGetDB() which uses getMaintenanceConnectionRef
 		$mockLoadBalancer->expects( $this->atLeastOnce() )
-			->method( 'getConnection' )
+			->method( 'getMaintenanceConnectionRef' )
-			->willReturn( $db );
+			->willReturnCallback( $this->getMockDBConnRefCallback( $mockLoadBalancer, $db ) );
 
 		// First call to Title::newFromID, faking no result (db lag?)
 		$db->expects( $this->at( 0 ) )
@@ -192,15 +209,15 @@ class RevisionStoreTest extends MediaWikiTestCase {
 		$this->setService( 'DBLoadBalancer', $mockLoadBalancer );
 
 		$db = $this->getMockDatabase();
-		// Title calls wfGetDB() which uses a regular Connection
+		// Title calls wfGetDB() which uses getMaintenanceConnectionRef
 		// Assert that the first call uses a REPLICA and the second falls back to master
-		$mockLoadBalancer->expects( $this->exactly( 2 ) )
-			->method( 'getConnection' )
-			->willReturn( $db );
-		// RevisionStore getTitle uses a ConnectionRef
 		$mockLoadBalancer->expects( $this->atLeastOnce() )
 			->method( 'getConnectionRef' )
-			->willReturn( $db );
+			->willReturnCallback( $this->getMockDBConnRefCallback( $mockLoadBalancer, $db ) );
+		// Title calls wfGetDB() which uses getMaintenanceConnectionRef
+		$mockLoadBalancer->expects( $this->exactly( 2 ) )
+			->method( 'getMaintenanceConnectionRef' )
+			->willReturnCallback( $this->getMockDBConnRefCallback( $mockLoadBalancer, $db ) );
 
 		// First call to Title::newFromID, faking no result (db lag?)
 		$db->expects( $this->at( 0 ) )
@@ -251,14 +268,14 @@ class RevisionStoreTest extends MediaWikiTestCase {
 		$this->setService( 'DBLoadBalancer', $mockLoadBalancer );
 
 		$db = $this->getMockDatabase();
-		// Title calls wfGetDB() which uses a regular Connection
-		$mockLoadBalancer->expects( $this->atLeastOnce() )
-			->method( 'getConnection' )
-			->willReturn( $db );
-		// RevisionStore getTitle uses a ConnectionRef
 		$mockLoadBalancer->expects( $this->atLeastOnce() )
 			->method( 'getConnectionRef' )
-			->willReturn( $db );
+			->willReturnCallback( $this->getMockDBConnRefCallback( $mockLoadBalancer, $db ) );
+		// Title calls wfGetDB() which uses getMaintenanceConnectionRef
+		// RevisionStore getTitle uses getMaintenanceConnectionRef
+		$mockLoadBalancer->expects( $this->atLeastOnce() )
+			->method( 'getMaintenanceConnectionRef' )
+			->willReturnCallback( $this->getMockDBConnRefCallback( $mockLoadBalancer, $db ) );
 
 		// First call to Title::newFromID, faking no result (db lag?)
 		$db->expects( $this->at( 0 ) )
@@ -299,15 +316,15 @@ class RevisionStoreTest extends MediaWikiTestCase {
 		$this->setService( 'DBLoadBalancer', $mockLoadBalancer );
 
 		$db = $this->getMockDatabase();
-		// Title calls wfGetDB() which uses a regular Connection
 		// Assert that the first call uses a REPLICA and the second falls back to master
-		$mockLoadBalancer->expects( $this->exactly( 2 ) )
+		// RevisionStore uses getMaintenanceConnectionRef
-			->method( 'getConnection' )
-			->willReturn( $db );
-		// RevisionStore getTitle uses a ConnectionRef
 		$mockLoadBalancer->expects( $this->atLeastOnce() )
 			->method( 'getConnectionRef' )
-			->willReturn( $db );
+			->willReturnCallback( $this->getMockDBConnRefCallback( $mockLoadBalancer, $db ) );
+		// Title calls wfGetDB() which uses getMaintenanceConnectionRef
+		$mockLoadBalancer->expects( $this->exactly( 2 ) )
+			->method( 'getMaintenanceConnectionRef' )
+			->willReturnCallback( $this->getMockDBConnRefCallback( $mockLoadBalancer, $db ) );
 
 		// First call to Title::newFromID, faking no result (db lag?)
 		$db->expects( $this->at( 0 ) )
@@ -368,12 +385,14 @@ class RevisionStoreTest extends MediaWikiTestCase {
 		$this->setService( 'DBLoadBalancer', $mockLoadBalancer );
 
 		$db = $this->getMockDatabase();
-		// Title calls wfGetDB() which uses a regular Connection
+		// Title calls wfGetDB() which uses getMaintenanceConnectionRef
 		// Assert that the first call uses a REPLICA and the second falls back to master
 
 		// RevisionStore getTitle uses getConnectionRef
-		// Title::newFromID uses getConnection
+		// Title::newFromID uses getMaintenanceConnectionRef
-		foreach ( [ 'getConnection', 'getConnectionRef' ] as $method ) {
+		foreach ( [
+			'getConnectionRef', 'getMaintenanceConnectionRef'
+		] as $method ) {
 			$mockLoadBalancer->expects( $this->exactly( 2 ) )
 				->method( $method )
 				->willReturnCallback( function ( $masterOrReplica ) use ( $db ) {