1 files changed, 36 insertions, 13 deletions

diff --git a/MLEB/Babel/includes/Database.php b/MLEB/Babel/includes/Database.php
index 5c0c015f..8dbb67f2 100644
--- a/MLEB/Babel/includes/Database.php
+++ b/MLEB/Babel/includes/Database.php
@@ -22,25 +22,27 @@ namespace MediaWiki\Babel;
 
 use MediaWiki\MediaWikiServices;
 use Wikimedia\Rdbms\IDatabase;
-use Wikimedia\Rdbms\LoadBalancer;
+use Wikimedia\Rdbms\LBFactory;
 
 class Database {
 
 	/**
-	 * @var LoadBalancer
+	 * @var LBFactory
 	 */
-	private $loadBalancer;
+	private $loadBalancerFactory;
 
 	public function __construct() {
-		$this->loadBalancer = MediaWikiServices::getInstance()->getDBLoadBalancer();
+		$this->loadBalancerFactory = MediaWikiServices::getInstance()->getDBLoadBalancerFactory();
 	}
 
 	/**
 	 * @param int $index
+	 * @param string|bool $wiki Database name if querying a different wiki
 	 * @return IDatabase
 	 */
-	protected function getDB( $index ) {
-		return $this->loadBalancer->getLazyConnectionRef( $index );
+	protected function getDB( $index, $wiki = false ) {
+		return $this->loadBalancerFactory->getMainLB( $wiki )
+			->getLazyConnectionRef( $index, [], $wiki );
 	}
 
 	/**
@@ -64,18 +66,39 @@ class Database {
 	}
 
 	/**
-	 * @param string $id
+	 * @param string $wiki Database name
+	 * @param string $username
+	 * @return string[] [ lang => level ]
+	 */
+	public function getForRemoteUser( $wiki, $username ) {
+		$rows = $this->getDB( DB_REPLICA, $wiki )->select(
+			[ 'babel', 'user' ],
+			[ 'babel_lang', 'babel_level' ],
+			[
+				'user_name' => $username
+			],
+			__METHOD__,
+			[],
+			[
+				'user' => [ 'INNER JOIN', 'babel_user=user_id' ]
+			]
+		);
+
+		$return = [];
+		foreach ( $rows as $row ) {
+			$return[$row->babel_lang] = $row->babel_level;
+		}
+
+		return $return;
+	}
+
+	/**
+	 * @param int $id
 	 * @param string[] $data [ lang => level ]
 	 * @return bool true if changes to the db were made
 	 */
 	public function setForUser( $id, array $data ) {
 		$dbw = $this->getDB( DB_MASTER );
-		if ( !$dbw->tableExists( 'babel' ) ) {
-			// Populate data as long as the table exists, so
-			// we can pre-populate the database before switching
-			// reads to the database
-			return false;
-		}
 
 		$newRows = [];
 		foreach ( $data as $lang => $level ) {