rdbms: Replace getConnection with getLazyConnectionRef

This would simplify the code for its users a lot.

Bug: T255493
Depends-On: I6ab4440a13f4682dd3a6e9cfc6c629222fbeab8a
Change-Id: I6e7544763bde56fc1e19de0358b71ba984a978ca

includes/externalstore/ExternalStoreDB.php

@@ -187,7 +187,7 @@ class ExternalStoreDB extends ExternalStoreMedium {
 	 * Get a primary database connection for the specified cluster
 	 *
 	 * @param string $cluster Cluster name
-	 * @return DBConnRef
+	 * @return \Wikimedia\Rdbms\IMaintainableDatabase
 	 * @since 1.37
 	 */
 	public function getPrimary( $cluster ) {
@@ -204,7 +204,7 @@ class ExternalStoreDB extends ExternalStoreMedium {
 	/**
 	 * @deprecated since 1.37; please use getPrimary() instead.
 	 * @param string $cluster Cluster name
-	 * @return DBConnRef
+	 * @return \Wikimedia\Rdbms\IMaintainableDatabase
 	 */
 	public function getMaster( $cluster ) {
 		wfDeprecated( __METHOD__, '1.37' );

includes/libs/rdbms/database/DBConnRef.php

@@ -62,7 +62,7 @@ class DBConnRef implements IMaintainableDatabase {
 	public function __call( $name, array $arguments ) {
 		if ( $this->conn === null ) {
 			list( $index, $groups, $wiki, $flags ) = $this->params;
-			$this->conn = $this->lb->getConnection( $index, $groups, $wiki, $flags );
+			$this->conn = $this->lb->getConnectionInternal( $index, $groups, $wiki, $flags );
 		}
 
 		return $this->conn->$name( ...$arguments );
@@ -109,12 +109,10 @@ class DBConnRef implements IMaintainableDatabase {
 			$domain = DatabaseDomain::newFromId( $this->params[self::FLD_DOMAIN] );
 			// Avoid triggering a database connection
 			return $domain->getTablePrefix();
-		} elseif ( $this->conn !== null && $prefix === null ) {
+		} else {
 			// This will just return the prefix
 			return $this->__call( __FUNCTION__, func_get_args() );
 		}
-		// Disallow things that might confuse the LoadBalancer tracking
-		throw $this->getDomainChangeException();
 	}
 
 	public function dbSchema( $schema = null ) {
@@ -850,7 +848,7 @@ class DBConnRef implements IMaintainableDatabase {
 	 */
 	public function __destruct() {
 		if ( $this->conn ) {
-			$this->lb->reuseConnection( $this->conn );
+			$this->lb->reuseConnectionInternal( $this->conn );
 		}
 	}
 }

includes/libs/rdbms/loadbalancer/ILoadBalancer.php

@@ -23,7 +23,6 @@
 namespace Wikimedia\Rdbms;
 
 use InvalidArgumentException;
-use LogicException;
 
 /**
  * Database cluster connection, tracking, load balancing, and transaction manager interface
@@ -219,7 +218,7 @@ interface ILoadBalancer {
 	public function getAnyOpenConnection( $i, $flags = 0 );
 
 	/**
-	 * Get a live handle for a specific or virtual (DB_PRIMARY/DB_REPLICA) server index
+	 * Get a lazy handle for a specific or virtual (DB_PRIMARY/DB_REPLICA) server index
 	 *
 	 * The server index, $i, can be one of the following:
 	 *   - DB_REPLICA: a server index will be selected by the load balancer based on read
@@ -237,9 +236,7 @@ interface ILoadBalancer {
 	 *      server selection method is usually only useful for internal load balancing logic.
 	 *      The value of $groups should be [] when using a specific server index.
 	 *
-	 * Handles acquired by this method, getConnectionRef(), getLazyConnectionRef(), and
+	 * Callers that get a *local* DB domain handle for the same server will share one handle for all of those
-	 * getMaintenanceConnectionRef() use the same set of shared connection pools. Callers that
-	 * get a *local* DB domain handle for the same server will share one handle for all of those
 	 * callers using CONN_TRX_AUTOCOMMIT (via $flags) and one handle for all of those callers not
 	 * using CONN_TRX_AUTOCOMMIT. Callers that get a *foreign* DB domain handle (via $domain) will
 	 * share any handle that has the right CONN_TRX_AUTOCOMMIT mode and is already on the right
@@ -253,10 +250,6 @@ interface ILoadBalancer {
 	 * query grouped (the default), DB_REPLICA handles. All such callers will operate within a
 	 * single database transaction as a consequence.
 	 *
-	 * Callers of this function that use a non-local $domain must call reuseConnection() after
-	 * their last query on this handle executed. This lets the load balancer share the handle with
-	 * other callers requesting connections on different database domains.
-	 *
 	 * Use CONN_TRX_AUTOCOMMIT to use a separate pool of only auto-commit handles. This flag
 	 * is ignored for databases with ATTR_DB_LEVEL_LOCKING (e.g. sqlite) in order to avoid
 	 * deadlocks. getServerAttributes() can be used to check such attributes beforehand. Avoid
@@ -302,30 +295,19 @@ interface ILoadBalancer {
 	public function getServerConnection( $i, $domain, $flags = 0 );
 
 	/**
-	 * Mark a live handle as being available for reuse under a different database domain
+	 * @deprecated since 1.39 noop
-	 *
-	 * This mechanism is reference-counted, and must be called the same number of times as
-	 * getConnection() to work. Never call this on handles acquired via getConnectionRef(),
-	 * getLazyConnectionRef(), and getMaintenanceConnectionRef(), as they already manage
-	 * the logic of calling this method when they fall out of scope in PHP.
-	 *
-	 * @see ILoadBalancer::getConnection()
-	 *
 	 * @param IDatabase $conn
-	 * @throws LogicException
 	 */
 	public function reuseConnection( IDatabase $conn );
 
 	/**
-	 * Get a live database handle reference for a server index
+	 * @internal Only to be used by DBConnRef
-	 *
+	 * @param IDatabase $conn
-	 * The CONN_TRX_AUTOCOMMIT flag is ignored for databases with ATTR_DB_LEVEL_LOCKING
+	 */
-	 * (e.g. sqlite) in order to avoid deadlocks. The getServerAttributes() method can be used
+	public function reuseConnectionInternal( IDatabase $conn );
-	 * to check such flags beforehand. Avoid the use of begin() or startAtomic()
+
-	 * on any CONN_TRX_AUTOCOMMIT connections.
+	/**
-	 *
+	 * @deprecated since 1.39, use ILoadBalancer::getConnection() instead.
-	 * @see ILoadBalancer::getConnection() for parameter information
-	 *
 	 * @param int $i Specific or virtual (DB_PRIMARY/DB_REPLICA) server index
 	 * @param string[]|string $groups Query group(s) in preference order; [] for the default group
 	 * @param string|bool $domain DB domain ID or false for the local domain
@@ -334,6 +316,16 @@ interface ILoadBalancer {
 	 */
 	public function getConnectionRef( $i, $groups = [], $domain = false, $flags = 0 ): IDatabase;
 
+	/**
+	 * @internal Only to be used by DBConnRef
+	 * @param int $i Specific or virtual (DB_PRIMARY/DB_REPLICA) server index
+	 * @param string[]|string $groups Query group(s) in preference order; [] for the default group
+	 * @param string|bool $domain DB domain ID or false for the local domain
+	 * @param int $flags Bitfield of CONN_* class constants (e.g. CONN_TRX_AUTOCOMMIT)
+	 * @return IDatabase
+	 */
+	public function getConnectionInternal( $i, $groups = [], $domain = false, $flags = 0 ): IDatabase;
+
 	/**
 	 * Get a lazy-connecting database handle reference for a server index
 	 *
@@ -346,7 +338,7 @@ interface ILoadBalancer {
 	 * on any CONN_TRX_AUTOCOMMIT connections.
 	 *
 	 * @see ILoadBalancer::getConnection() for parameter information
-	 *
+	 * @deprecated since 1.38, use ILoadBalancer::getConnectionRef() instead.
 	 * @param int $i Specific or virtual (DB_PRIMARY/DB_REPLICA) server index
 	 * @param string[]|string $groups Query group(s) in preference order; [] for the default group
 	 * @param string|bool $domain DB domain ID or false for the local domain
@@ -354,7 +346,6 @@ interface ILoadBalancer {
 	 * @return IDatabase Live connection handle
 	 * @throws DBError If no live handle could be obtained
 	 * @throws DBAccessError If disable() was previously called
-	 * @deprecated since 1.38 use getConnectionRef instead
 	 */
 	public function getLazyConnectionRef( $i, $groups = [], $domain = false, $flags = 0 ): IDatabase;
 

includes/libs/rdbms/loadbalancer/LoadBalancer.php

@@ -670,7 +670,7 @@ class LoadBalancer implements ILoadBalancerForOwner {
 			// Decrement reference counter, we are finished with this connection.
 			// It will be incremented for the caller later.
 			if ( !$this->localDomain->equals( $domain ) ) {
-				$this->reuseConnection( $conn );
+				$this->reuseConnectionInternal( $conn );
 			}
 
 			// Return this server
@@ -907,6 +907,18 @@ class LoadBalancer implements ILoadBalancerForOwner {
 	}
 
 	public function getConnection( $i, $groups = [], $domain = false, $flags = 0 ) {
+		return $this->getConnectionRef( $i, $groups, $domain, $flags );
+	}
+
+	/**
+	 * @internal
+	 * @param int $i Specific (overrides $groups) or virtual (DB_PRIMARY/DB_REPLICA) server index
+	 * @param string[]|string $groups Query group(s) in preference order; [] for the default group
+	 * @param string|bool $domain DB domain ID or false for the local domain
+	 * @param int $flags Bitfield of CONN_* class constants
+	 * @return IDatabase
+	 */
+	public function getConnectionInternal( $i, $groups = [], $domain = false, $flags = 0 ): IDatabase {
 		$domain = $this->resolveDomainID( $domain );
 		$groups = $this->resolveGroups( $groups, $i );
 		$flags = $this->sanitizeConnectionFlags( $flags, $i, $domain );
@@ -930,7 +942,6 @@ class LoadBalancer implements ILoadBalancerForOwner {
 				: 'The database is read-only until replica database servers becomes reachable.';
 			$conn->setLBInfo( $conn::LB_READ_ONLY_REASON, $reason );
 		}
-
 		return $conn;
 	}
 
@@ -991,6 +1002,10 @@ class LoadBalancer implements ILoadBalancerForOwner {
 	}
 
 	public function reuseConnection( IDatabase $conn ) {
+		// no-op
+	}
+
+	public function reuseConnectionInternal( IDatabase $conn ) {
 		$serverIndex = $conn->getLBInfo( self::INFO_SERVER_INDEX );
 		$refCount = $conn->getLBInfo( self::INFO_FOREIGN_REF_COUNT );
 		if ( $serverIndex === null || $refCount === null ) {
@@ -1074,7 +1089,7 @@ class LoadBalancer implements ILoadBalancerForOwner {
 
 		$domain = $this->resolveDomainID( $domain );
 		$role = $this->getRoleFromIndex( $i );
-		$conn = $this->getConnection( $i, $groups, $domain, $flags );
+		$conn = $this->getConnectionInternal( $i, $groups, $domain, $flags );
 
 		return new DBConnRef( $this, $conn, $role );
 	}
@@ -2133,7 +2148,7 @@ class LoadBalancer implements ILoadBalancerForOwner {
 					} catch ( DBError $e ) {
 						$readOnly = 0;
 					}
-					$this->reuseConnection( $conn );
+					$this->reuseConnectionInternal( $conn );
 				} else {
 					$readOnly = 0;
 				}

tests/phpunit/MediaWikiIntegrationTestCase.php

@@ -434,7 +434,7 @@ abstract class MediaWikiIntegrationTestCase extends PHPUnit\Framework\TestCase {
 				$useTemporaryTables = !$this->getCliArg( 'use-normal-tables' );
 
 				$lb = MediaWikiServices::getInstance()->getDBLoadBalancer();
-				$this->db = $lb->getConnection( DB_PRIMARY );
+				$this->db = $lb->getConnectionInternal( DB_PRIMARY );
 
 				$this->checkDbIsSupported();
 

tests/phpunit/includes/db/LBFactoryTest.php

@@ -479,10 +479,10 @@ class LBFactoryTest extends MediaWikiIntegrationTestCase {
 			"Correct full table name"
 		);
 
-		$lb->reuseConnection( $db ); // don't care
+		unset( $db );
 
-		$db = $lb->getConnection( DB_PRIMARY ); // local domain connection
 		$factory->setLocalDomainPrefix( 'my_' );
+		$db = $lb->getConnection( DB_PRIMARY ); // local domain connection
 
 		$this->assertEquals( $wgDBname, $db->getDBname() );
 		$this->assertEquals(
@@ -548,7 +548,7 @@ class LBFactoryTest extends MediaWikiIntegrationTestCase {
 			"Correct full table name"
 		);
 
-		$lb->reuseConnection( $db ); // don't care
+		unset( $db );
 
 		$factory->setLocalDomainPrefix( 'my_' );
 		$db = $lb->getConnection( DB_PRIMARY, [], "$wgDBname-my_" );
@@ -596,12 +596,8 @@ class LBFactoryTest extends MediaWikiIntegrationTestCase {
 		/** @var IDatabase $db */
 		$db = $lb->getConnection( DB_PRIMARY, [], $lb::DOMAIN_ANY );
 
-		try {
+		$this->expectException( \Wikimedia\Rdbms\DBUnexpectedError::class );
-			$this->assertFalse( @$db->selectDomain( 'garbagedb' ) );
+		$db->selectDomain( 'garbagedb' );
-			$this->fail( "No error thrown." );
-		} catch ( \Wikimedia\Rdbms\DBQueryError $e ) {
-			$this->assertRegExp( '/[\'"]garbagedb[\'"]/', $e->getMessage() );
-		}
 	}
 
 	/**
@@ -626,7 +622,7 @@ class LBFactoryTest extends MediaWikiIntegrationTestCase {
 		}
 
 		/** @var IDatabase $db */
-		$this->assertNotNull( $lb->getConnection( DB_PRIMARY, [], $lb::DOMAIN_ANY ) );
+		$this->assertNotNull( $lb->getConnectionInternal( DB_PRIMARY, [], $lb::DOMAIN_ANY ) );
 	}
 
 	/**
@@ -650,7 +646,7 @@ class LBFactoryTest extends MediaWikiIntegrationTestCase {
 			$this->markTestSkipped( "Not applicable per databasesAreIndependent()" );
 		}
 
-		$db = $lb->getConnection( DB_PRIMARY );
+		$db = $lb->getConnectionInternal( DB_PRIMARY );
 		$db->selectDomain( 'garbage-db' );
 	}
 

tests/phpunit/includes/db/LoadBalancerTest.php

@@ -96,6 +96,7 @@ class LoadBalancerTest extends MediaWikiIntegrationTestCase {
 		$this->assertFalse( $called );
 
 		$dbw = $lb->getConnection( DB_PRIMARY );
+		$dbw->getServerName();
 		$this->assertTrue( $called );
 		$this->assertEquals(
 			$dbw::ROLE_STREAMING_MASTER, $dbw->getTopologyRole(), 'master shows as master'
@@ -108,7 +109,7 @@ class LoadBalancerTest extends MediaWikiIntegrationTestCase {
 			$dbr::ROLE_STREAMING_MASTER, $dbr->getTopologyRole(), 'DB_REPLICA also gets the master' );
 		$this->assertTrue( $dbr->getFlag( $dbw::DBO_TRX ), "DBO_TRX set on replica" );
 
-		if ( !$lb->getServerAttributes( $lb->getWriterIndex() )[$dbw::ATTR_DB_LEVEL_LOCKING] ) {
+		if ( !$lb->getServerAttributes( $lb->getWriterIndex() )[Database::ATTR_DB_LEVEL_LOCKING] ) {
 			$dbwAuto = $lb->getConnection( DB_PRIMARY, [], false, $lb::CONN_TRX_AUTOCOMMIT );
 			$this->assertFalse(
 				$dbwAuto->getFlag( $dbw::DBO_TRX ), "No DBO_TRX with CONN_TRX_AUTOCOMMIT" );
@@ -124,6 +125,7 @@ class LoadBalancerTest extends MediaWikiIntegrationTestCase {
 				$dbr, $dbrAuto, "CONN_TRX_AUTOCOMMIT uses separate connection" );
 
 			$dbwAuto2 = $lb->getConnection( DB_PRIMARY, [], false, $lb::CONN_TRX_AUTOCOMMIT );
+			$dbwAuto2->getServerName();
 			$this->assertEquals( $dbwAuto2, $dbwAuto, "CONN_TRX_AUTOCOMMIT reuses connections" );
 		}
 
@@ -188,7 +190,7 @@ class LoadBalancerTest extends MediaWikiIntegrationTestCase {
 		$this->assertTrue( $dbr->getFlag( $dbw::DBO_TRX ), "DBO_TRX set on replica" );
 		$this->assertEquals( $dbr->getLBInfo( 'serverIndex' ), $lb->getReaderIndex() );
 
-		if ( !$lb->getServerAttributes( $lb->getWriterIndex() )[$dbw::ATTR_DB_LEVEL_LOCKING] ) {
+		if ( !$lb->getServerAttributes( $lb->getWriterIndex() )[Database::ATTR_DB_LEVEL_LOCKING] ) {
 			$dbwAuto = $lb->getConnection( DB_PRIMARY, [], false, $lb::CONN_TRX_AUTOCOMMIT );
 			$this->assertFalse(
 				$dbwAuto->getFlag( $dbw::DBO_TRX ), "No DBO_TRX with CONN_TRX_AUTOCOMMIT" );
@@ -204,6 +206,7 @@ class LoadBalancerTest extends MediaWikiIntegrationTestCase {
 				$dbr, $dbrAuto, "CONN_TRX_AUTOCOMMIT uses separate connection" );
 
 			$dbwAuto2 = $lb->getConnection( DB_PRIMARY, [], false, $lb::CONN_TRX_AUTOCOMMIT );
+			$dbwAuto2->getServerName();
 			$this->assertEquals( $dbwAuto2, $dbwAuto, "CONN_TRX_AUTOCOMMIT reuses connections" );
 		}
 
@@ -354,12 +357,13 @@ class LoadBalancerTest extends MediaWikiIntegrationTestCase {
 		}
 	}
 
-	private function assertWriteAllowed( Database $db ) {
+	private function assertWriteAllowed( IMaintainableDatabase $db ) {
 		$table = $db->tableName( 'some_table' );
 		// Trigger a transaction so that rollback() will remove all the tables.
 		// Don't do this for MySQL as it auto-commits transactions for DDL
 		// statements such as CREATE TABLE.
 		$useAtomicSection = in_array( $db->getType(), [ 'sqlite', 'postgres' ], true );
+		/** @var Database $db */
 		try {
 			$db->dropTable( 'some_table' );
 			$this->assertNotEquals( TransactionManager::STATUS_TRX_ERROR, $db->trxStatus() );
@@ -442,8 +446,7 @@ class LoadBalancerTest extends MediaWikiIntegrationTestCase {
 	}
 
 	/**
-	 * @covers \Wikimedia\Rdbms\LoadBalancer::getConnection()
+	 * @covers \Wikimedia\Rdbms\LoadBalancer::getConnectionInternal()
-	 * @covers \Wikimedia\Rdbms\LoadBalancer::getAnyOpenConnection()
 	 * @covers \Wikimedia\Rdbms\LoadBalancer::getWriterIndex()
 	 */
 	public function testOpenConnection() {
@@ -452,12 +455,12 @@ class LoadBalancerTest extends MediaWikiIntegrationTestCase {
 		$i = $lb->getWriterIndex();
 		$this->assertFalse( $lb->getAnyOpenConnection( $i ) );
 
-		$conn1 = $lb->getConnection( $i );
+		$conn1 = $lb->getConnectionInternal( $i );
+		$conn1->getServerName();
 		$this->assertNotEquals( null, $conn1 );
-		$this->assertEquals( $conn1, $lb->getAnyOpenConnection( $i ) );
 		$this->assertFalse( $conn1->getFlag( DBO_TRX ) );
 
-		$conn2 = $lb->getConnection( $i, [], false, $lb::CONN_TRX_AUTOCOMMIT );
+		$conn2 = $lb->getConnectionInternal( $i, [], false, $lb::CONN_TRX_AUTOCOMMIT );
 		$this->assertNotEquals( null, $conn2 );
 		$this->assertFalse( $conn2->getFlag( DBO_TRX ) );
 
@@ -465,26 +468,8 @@ class LoadBalancerTest extends MediaWikiIntegrationTestCase {
 			$this->assertFalse(
 				$lb->getAnyOpenConnection( $i, $lb::CONN_TRX_AUTOCOMMIT ) );
 			$this->assertEquals( $conn1,
-				$lb->getConnection(
+				$lb->getConnectionInternal(
 					$i, [], false, $lb::CONN_TRX_AUTOCOMMIT ), $lb::CONN_TRX_AUTOCOMMIT );
-		} else {
-			$this->assertEquals( $conn2,
-				$lb->getAnyOpenConnection( $i, $lb::CONN_TRX_AUTOCOMMIT ) );
-			$this->assertEquals( $conn2,
-				$lb->getConnection( $i, [], false, $lb::CONN_TRX_AUTOCOMMIT ) );
-
-			$conn2->startAtomic( __METHOD__ );
-			try {
-				$lb->getConnection( $i, [], false, $lb::CONN_TRX_AUTOCOMMIT );
-				$conn2->endAtomic( __METHOD__ );
-				$this->fail( "No exception thrown." );
-			} catch ( DBUnexpectedError $e ) {
-				$this->assertEquals(
-					'Handle requested with CONN_TRX_AUTOCOMMIT yet it has a transaction',
-					$e->getMessage()
-				);
-			}
-			$conn2->endAtomic( __METHOD__ );
 		}
 
 		$lb->closeAll( __METHOD__ );
@@ -525,7 +510,9 @@ class LoadBalancerTest extends MediaWikiIntegrationTestCase {
 		] );
 
 		$conn1 = $lb->getConnection( $lb->getWriterIndex(), [], false );
+		$conn1->getServerName();
 		$conn2 = $lb->getConnection( $lb->getWriterIndex(), [], '' );
+		$conn2->getServerName();
 
 		/** @var LoadBalancer $lbWrapper */
 		$lbWrapper = TestingAccessWrapper::newFromObject( $lb );
@@ -588,9 +575,6 @@ class LoadBalancerTest extends MediaWikiIntegrationTestCase {
 		$lb->flushPrimarySessions( __METHOD__ );
 		$this->assertSame( TransactionManager::STATUS_TRX_NONE, $conn1->trxStatus() );
 		$this->assertSame( TransactionManager::STATUS_TRX_NONE, $conn2->trxStatus() );
-
-		$conn1->close();
-		$conn2->close();
 	}
 
 	/**

tests/phpunit/maintenance/backup_PageTest.php

@@ -61,7 +61,7 @@ class BackupDumperPageTest extends DumpTestCase {
 			$this->streamingLoadBalancer = $lbFactory->newMainLB();
 		}
 
-		$db = $this->streamingLoadBalancer->getConnection( DB_REPLICA );
+		$db = $this->streamingLoadBalancer->getConnection( DB_PRIMARY );
 
 		// Make sure the DB connection has the fake table clones and the fake table prefix
 		$this->dbClone = MediaWikiIntegrationTestCase::setupDatabaseWithTestPrefix( $db );

tests/phpunit/tests/MediaWikiIntegrationTestCaseTest.php

@@ -222,7 +222,7 @@ class MediaWikiIntegrationTestCaseTest extends MediaWikiIntegrationTestCase {
 
 		$lbFactory = $this->getServiceContainer()->getDBLoadBalancerFactory();
 		$lb = $lbFactory->newMainLB();
-		$db = $lb->getConnection( DB_REPLICA );
+		$db = $lb->getConnection( DB_PRIMARY );
 
 		$this->assertNotSame( $this->db, $db );
 

tests/phpunit/unit/includes/libs/rdbms/database/DBConnRefTest.php

@@ -25,6 +25,12 @@ class DBConnRefTest extends PHPUnit\Framework\TestCase {
 			}
 		);
 
+		$lb->method( 'getConnectionInternal' )->willReturnCallback(
+			function () {
+				return $this->getDatabaseMock();
+			}
+		);
+
 		$lb->method( 'getConnectionRef' )->willReturnCallback(
 			function () use ( $lb ) {
 				return $this->getDBConnRef( $lb );
@@ -82,7 +88,7 @@ class DBConnRefTest extends PHPUnit\Framework\TestCase {
 		$lb = $this->createMock( ILoadBalancer::class );
 
 		$lb->expects( $this->once() )
-			->method( 'getConnection' )
+			->method( 'getConnectionInternal' )
 			->with( DB_PRIMARY, [ 'test' ], 'dummy', ILoadBalancer::CONN_TRX_AUTOCOMMIT )
 			->willReturnCallback(
 				function () {
@@ -111,7 +117,7 @@ class DBConnRefTest extends PHPUnit\Framework\TestCase {
 		$lb = $this->getLoadBalancerMock();
 
 		$lb->expects( $this->once() )
-			->method( 'reuseConnection' );
+			->method( 'reuseConnectionInternal' );
 
 		$this->innerMethodForTestDestruct( $lb );
 	}

tests/phpunit/unit/includes/libs/rdbms/database/DatabaseTest.php

@@ -242,7 +242,7 @@ class DatabaseTest extends PHPUnit\Framework\TestCase {
 		// Ask for the connection so that LB sets internal state
 		// about this connection being the primary connection
 		$lb = $lbFactory->getMainLB();
-		$conn = $lb->getConnection( $lb->getWriterIndex() );
+		$conn = $lb->getConnectionInternal( $lb->getWriterIndex() );
 		$this->assertSame( $db, $conn, 'Same DB instance' );
 		$this->assertTrue( $db->getFlag( DBO_TRX ), 'DBO_TRX is set' );
 
@@ -339,7 +339,7 @@ class DatabaseTest extends PHPUnit\Framework\TestCase {
 		// Ask for the connection so that LB sets internal state
 		// about this connection being the primary connection
 		$lb = $lbFactory->getMainLB();
-		$conn = $lb->getConnection( $lb->getWriterIndex() );
+		$conn = $lb->getConnectionInternal( $lb->getWriterIndex() );
 		$this->assertSame( $db, $conn, 'Same DB instance' );
 
 		$this->assertFalse( $lb->hasPrimaryChanges() );