From ef381994222efcc1ce9a380d340b9cb51ad458ad Mon Sep 17 00:00:00 2001 From: Aaron Schulz Date: Sun, 14 Jul 2019 19:54:47 -0700 Subject: [PATCH] rdbms: fix IDatabase::setLBInfo() handling of null and allow clearing keys Change-Id: I20cb799b54cabb1172940f8ece93b7f45d7cf0ba --- includes/libs/rdbms/database/DBConnRef.php | 2 +- includes/libs/rdbms/database/Database.php | 14 +++++++--- includes/libs/rdbms/database/IDatabase.php | 16 +++++------ .../libs/rdbms/loadbalancer/LoadBalancer.php | 2 +- .../libs/rdbms/database/DatabaseTest.php | 27 +++++++++++++++++++ 5 files changed, 46 insertions(+), 15 deletions(-) diff --git a/includes/libs/rdbms/database/DBConnRef.php b/includes/libs/rdbms/database/DBConnRef.php index 2cc2c90d69..2c9858add5 100644 --- a/includes/libs/rdbms/database/DBConnRef.php +++ b/includes/libs/rdbms/database/DBConnRef.php @@ -130,7 +130,7 @@ class DBConnRef implements IDatabase { return $this->__call( __FUNCTION__, func_get_args() ); } - public function setLBInfo( $name, $value = null ) { + public function setLBInfo( $nameOrArray, $value = null ) { // Disallow things that might confuse the LoadBalancer tracking throw new DBUnexpectedError( $this, "Changing LB info is disallowed to enable reuse." ); } diff --git a/includes/libs/rdbms/database/Database.php b/includes/libs/rdbms/database/Database.php index 3024b00e98..60062fbc13 100644 --- a/includes/libs/rdbms/database/Database.php +++ b/includes/libs/rdbms/database/Database.php @@ -586,11 +586,17 @@ abstract class Database implements IDatabase, IMaintainableDatabase, LoggerAware return null; } - public function setLBInfo( $name, $value = null ) { - if ( is_null( $value ) ) { - $this->lbInfo = $name; + public function setLBInfo( $nameOrArray, $value = null ) { + if ( is_array( $nameOrArray ) ) { + $this->lbInfo = $nameOrArray; + } elseif ( is_string( $nameOrArray ) ) { + if ( $value !== null ) { + $this->lbInfo[$nameOrArray] = $value; + } else { + unset( $this->lbInfo[$nameOrArray] ); + } } else { - $this->lbInfo[$name] = $value; + throw new InvalidArgumentException( "Got non-string key" ); } } diff --git a/includes/libs/rdbms/database/IDatabase.php b/includes/libs/rdbms/database/IDatabase.php index 3b9d1afa4f..ef7f1e24f6 100644 --- a/includes/libs/rdbms/database/IDatabase.php +++ b/includes/libs/rdbms/database/IDatabase.php @@ -223,14 +223,12 @@ interface IDatabase { public function getLBInfo( $name = null ); /** - * Set the LB info array, or a member of it. If called with one parameter, - * the LB info array is set to that parameter. If it is called with two - * parameters, the member with the given name is set to the given value. + * Set the entire array or a particular key of the managing load balancer info array * - * @param array|string $name - * @param array|null $value + * @param array|string $nameOrArray The new array or the name of a key to set + * @param array|null $value If $nameOrArray is a string, the new key value (null to unset) */ - public function setLBInfo( $name, $value = null ); + public function setLBInfo( $nameOrArray, $value = null ); /** * Set a lazy-connecting DB handle to the master DB (for replication status purposes) @@ -1158,7 +1156,7 @@ interface IDatabase { /** * Change the current database * - * This should not be called outside LoadBalancer for connections managed by a LoadBalancer + * This should only be called by a load balancer or if the handle is not attached to one * * @param string $db * @return bool True unless an exception was thrown @@ -1171,9 +1169,9 @@ interface IDatabase { /** * Set the current domain (database, schema, and table prefix) * - * This will throw an error for some database types if the database unspecified + * This will throw an error for some database types if the database is unspecified * - * This should not be called outside LoadBalancer for connections managed by a LoadBalancer + * This should only be called by a load balancer or if the handle is not attached to one * * @param string|DatabaseDomain $domain * @since 1.32 diff --git a/includes/libs/rdbms/loadbalancer/LoadBalancer.php b/includes/libs/rdbms/loadbalancer/LoadBalancer.php index cab0201ed4..7e94f80534 100644 --- a/includes/libs/rdbms/loadbalancer/LoadBalancer.php +++ b/includes/libs/rdbms/loadbalancer/LoadBalancer.php @@ -1861,7 +1861,7 @@ class LoadBalancer implements ILoadBalancer { } if ( $conn->getFlag( $conn::DBO_TRX ) ) { - $conn->setLBInfo( 'trxRoundId', false ); + $conn->setLBInfo( 'trxRoundId', null ); // remove the round ID } if ( $conn->getFlag( $conn::DBO_DEFAULT ) ) { diff --git a/tests/phpunit/includes/libs/rdbms/database/DatabaseTest.php b/tests/phpunit/includes/libs/rdbms/database/DatabaseTest.php index 8b24791ca6..482ab4b5f5 100644 --- a/tests/phpunit/includes/libs/rdbms/database/DatabaseTest.php +++ b/tests/phpunit/includes/libs/rdbms/database/DatabaseTest.php @@ -704,4 +704,31 @@ class DatabaseTest extends PHPUnit\Framework\TestCase { $this->assertSame( $oldDomain, $this->db->getDomainId() ); } + /** + * @covers Wikimedia\Rdbms\Database::getLBInfo + * @covers Wikimedia\Rdbms\Database::setLBInfo + */ + public function testGetSetLBInfo() { + $db = $this->getMockDB(); + + $this->assertEquals( [], $db->getLBInfo() ); + $this->assertNull( $db->getLBInfo( 'pringles' ) ); + + $db->setLBInfo( 'soda', 'water' ); + $this->assertEquals( [ 'soda' => 'water' ], $db->getLBInfo() ); + $this->assertNull( $db->getLBInfo( 'pringles' ) ); + $this->assertEquals( 'water', $db->getLBInfo( 'soda' ) ); + + $db->setLBInfo( 'basketball', 'Lebron' ); + $this->assertEquals( [ 'soda' => 'water', 'basketball' => 'Lebron' ], $db->getLBInfo() ); + $this->assertEquals( 'water', $db->getLBInfo( 'soda' ) ); + $this->assertEquals( 'Lebron', $db->getLBInfo( 'basketball' ) ); + + $db->setLBInfo( 'soda', null ); + $this->assertEquals( [ 'basketball' => 'Lebron' ], $db->getLBInfo() ); + + $db->setLBInfo( [ 'King' => 'James' ] ); + $this->assertNull( $db->getLBInfo( 'basketball' ) ); + $this->assertEquals( [ 'King' => 'James' ], $db->getLBInfo() ); + } } -- 2.20.1