rdbms: fix IDatabase::setLBInfo() handling of null and allow clearing keys

Change-Id: I20cb799b54cabb1172940f8ece93b7f45d7cf0ba

diff --git a/includes/libs/rdbms/database/DBConnRef.php b/includes/libs/rdbms/database/DBConnRef.php
index 2cc2c90..2c9858a 100644 (file)
--- 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 3024b00..60062fb 100644 (file)
--- 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 3b9d1af..ef7f1e2 100644 (file)
--- 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 cab0201..7e94f80 100644 (file)
--- 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 8b24791..482ab4b 100644 (file)
--- 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() );
+       }
 }