From: Aaron Schulz <aschulz@wikimedia.org>
Date: Mon, 19 Nov 2018 08:13:30 +0000 (-0800)
Subject: rdbms: add ILBFactory::redefineLocalDomain method
X-Git-Tag: 1.34.0-rc.0~3442^2
X-Git-Url: http://git.cyclocoop.org/%7B%24www_url%7Dadmin/compta/operations/?a=commitdiff_plain;h=1c205f1822b6214ac9c17a2e2affddafeecbac1b;p=lhc%2Fweb%2Fwiklou.git

rdbms: add ILBFactory::redefineLocalDomain method

This is intended for use with scripts like addWiki.php to avoid mismatched domain errors.

Bug: T209483
Change-Id: Ie24f83f4e0fad7b01690d2055e1529a4cb0275b5
---

diff --git a/includes/libs/rdbms/database/IDatabase.php b/includes/libs/rdbms/database/IDatabase.php
index 532d818771..27958838fb 100644
--- a/includes/libs/rdbms/database/IDatabase.php
+++ b/includes/libs/rdbms/database/IDatabase.php
@@ -1112,11 +1112,13 @@ interface IDatabase {
 	/**
 	 * Change the current database
 	 *
+	 * This should not be called outside LoadBalancer for connections managed by a LoadBalancer
+	 *
 	 * @param string $db
 	 * @return bool True unless an exception was thrown
 	 * @throws DBConnectionError If databasesAreIndependent() is true and an error occurs
 	 * @throws DBError
-	 * @deprecated Since 1.32
+	 * @deprecated Since 1.32 Use selectDomain() instead
 	 */
 	public function selectDB( $db );
 
@@ -1125,6 +1127,8 @@ interface IDatabase {
 	 *
 	 * This will throw an error for some database types if the database unspecified
 	 *
+	 * This should not be called outside LoadBalancer for connections managed by a LoadBalancer
+	 *
 	 * @param string|DatabaseDomain $domain
 	 * @since 1.32
 	 * @throws DBConnectionError
diff --git a/includes/libs/rdbms/lbfactory/ILBFactory.php b/includes/libs/rdbms/lbfactory/ILBFactory.php
index 57f60cafbb..8c95c165ca 100644
--- a/includes/libs/rdbms/lbfactory/ILBFactory.php
+++ b/includes/libs/rdbms/lbfactory/ILBFactory.php
@@ -83,6 +83,14 @@ interface ILBFactory {
 	 */
 	public function resolveDomainID( $domain );
 
+	/**
+	 * Close all connection and redefine the local domain for testing or schema creation
+	 *
+	 * @param DatabaseDomain|string $domain
+	 * @since 1.33
+	 */
+	public function redefineLocalDomain( $domain );
+
 	/**
 	 * Create a new load balancer object. The resulting object will be untracked,
 	 * not chronology-protected, and the caller is responsible for cleaning it up.
diff --git a/includes/libs/rdbms/lbfactory/LBFactory.php b/includes/libs/rdbms/lbfactory/LBFactory.php
index d213dc9e80..b60dd429cb 100644
--- a/includes/libs/rdbms/lbfactory/LBFactory.php
+++ b/includes/libs/rdbms/lbfactory/LBFactory.php
@@ -630,6 +630,16 @@ abstract class LBFactory implements ILBFactory {
 		} );
 	}
 
+	public function redefineLocalDomain( $domain ) {
+		$this->closeAll();
+
+		$this->localDomain = DatabaseDomain::newFromId( $domain );
+
+		$this->forEachLB( function ( ILoadBalancer $lb ) {
+			$lb->redefineLocalDomain( $this->localDomain );
+		} );
+	}
+
 	public function closeAll() {
 		$this->forEachLBCallMethod( 'closeAll', [] );
 	}
diff --git a/includes/libs/rdbms/loadbalancer/ILoadBalancer.php b/includes/libs/rdbms/loadbalancer/ILoadBalancer.php
index 78815f201a..661c184973 100644
--- a/includes/libs/rdbms/loadbalancer/ILoadBalancer.php
+++ b/includes/libs/rdbms/loadbalancer/ILoadBalancer.php
@@ -136,6 +136,14 @@ interface ILoadBalancer {
 	 */
 	public function resolveDomainID( $domain );
 
+	/**
+	 * Close all connection and redefine the local domain for testing or schema creation
+	 *
+	 * @param DatabaseDomain|string $domain
+	 * @since 1.33
+	 */
+	public function redefineLocalDomain( $domain );
+
 	/**
 	 * Get the index of the reader connection, which may be a replica DB
 	 *
diff --git a/includes/libs/rdbms/loadbalancer/LoadBalancer.php b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
index 7721707a54..a6858ef70c 100644
--- a/includes/libs/rdbms/loadbalancer/LoadBalancer.php
+++ b/includes/libs/rdbms/loadbalancer/LoadBalancer.php
@@ -1945,6 +1945,12 @@ class LoadBalancer implements ILoadBalancer {
 		} );
 	}
 
+	public function redefineLocalDomain( $domain ) {
+		$this->closeAll();
+
+		$this->setLocalDomain( DatabaseDomain::newFromId( $domain ) );
+	}
+
 	/**
 	 * @param DatabaseDomain $domain
 	 */
diff --git a/tests/phpunit/includes/db/LBFactoryTest.php b/tests/phpunit/includes/db/LBFactoryTest.php
index e865e79b68..ba02dd94e8 100644
--- a/tests/phpunit/includes/db/LBFactoryTest.php
+++ b/tests/phpunit/includes/db/LBFactoryTest.php
@@ -621,6 +621,50 @@ class LBFactoryTest extends MediaWikiTestCase {
 		}
 	}
 
+	public function testRedefineLocalDomain() {
+		global $wgDBname;
+
+		if ( wfGetDB( DB_MASTER )->databasesAreIndependent() ) {
+			self::markTestSkipped( "Skipping tests about selecting DBs: not applicable" );
+			return;
+		}
+
+		$factory = $this->newLBFactoryMulti(
+			[],
+			[]
+		);
+		$lb = $factory->getMainLB();
+
+		$conn1 = $lb->getConnectionRef( DB_MASTER );
+		$this->assertEquals(
+			wfWikiID(),
+			$conn1->getDomainID()
+		);
+		unset( $conn1 );
+
+		$factory->redefineLocalDomain( 'somedb-prefix' );
+		$this->assertEquals( 'somedb-prefix', $factory->getLocalDomainID() );
+
+		$domain = new DatabaseDomain( $wgDBname, null, 'pref' );
+		$factory->redefineLocalDomain( $domain );
+
+		$n = 0;
+		$lb->forEachOpenConnection( function () use ( &$n ) {
+			++$n;
+		} );
+		$this->assertEquals( 0, $n, "Connections closed" );
+
+		$conn2 = $lb->getConnectionRef( DB_MASTER );
+		$this->assertEquals(
+			$domain->getId(),
+			$conn2->getDomainID()
+		);
+		unset( $conn2 );
+
+		$factory->closeAll();
+		$factory->destroy();
+	}
+
 	private function quoteTable( Database $db, $table ) {
 		if ( $db->getType() === 'sqlite' ) {
 			return $table;