From 335066505a3ea6121f3c5b5eae43b0c37ca6e2eb Mon Sep 17 00:00:00 2001 From: Aaron Schulz Date: Sat, 6 Apr 2019 17:28:14 -0700 Subject: [PATCH] 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 | 4 +- .../includes/Revision/RevisionStoreTest.php | 67 ++++++++++++------- 2 files changed, 45 insertions(+), 26 deletions(-) diff --git a/includes/GlobalFunctions.php b/includes/GlobalFunctions.php index 7b4b502905..1741958681 100644 --- a/includes/GlobalFunctions.php +++ b/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 ); } /** diff --git a/tests/phpunit/includes/Revision/RevisionStoreTest.php b/tests/phpunit/includes/Revision/RevisionStoreTest.php index 0648bfce6e..83872e3a37 100644 --- a/tests/phpunit/includes/Revision/RevisionStoreTest.php +++ b/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' ) - ->willReturn( $db ); + ->method( 'getMaintenanceConnectionRef' ) + ->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 ) ) - ->method( 'getConnection' ) - ->willReturn( $db ); - // RevisionStore getTitle uses a ConnectionRef + // RevisionStore uses getMaintenanceConnectionRef $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 - foreach ( [ 'getConnection', 'getConnectionRef' ] as $method ) { + // Title::newFromID uses getMaintenanceConnectionRef + foreach ( [ + 'getConnectionRef', 'getMaintenanceConnectionRef' + ] as $method ) { $mockLoadBalancer->expects( $this->exactly( 2 ) ) ->method( $method ) ->willReturnCallback( function ( $masterOrReplica ) use ( $db ) { -- 2.20.1