From 51c17f193104220796c6fc6ce6c380c34242e682 Mon Sep 17 00:00:00 2001
From: Christian Aistleitner <christian@quelltextlich.at>
Date: Tue, 3 Apr 2012 11:23:20 +0200
Subject: [PATCH] Database dependency injection for BackupDumpers

Change-Id: I48f8cd176cc51fe133a8490e3e001a8be4978691
---
 maintenance/backup.inc         | 30 +++++++++++++++++++++++++++++-
 maintenance/backupTextPass.inc |  5 +++++
 2 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/maintenance/backup.inc b/maintenance/backup.inc
index 470a513f4e..6eccb26a6b 100644
--- a/maintenance/backup.inc
+++ b/maintenance/backup.inc
@@ -60,6 +60,15 @@ class BackupDumper {
 
 	var $outputTypes = array(), $filterTypes = array();
 
+	/**
+	 * The dependency-injected database to use.
+	 *
+	 * @var DatabaseBase|null
+	 *
+	 * @see self::setDb
+	 */
+	protected $forcedDb = null;
+
 	/**
 	 * @var LoadBalancer
 	 */
@@ -245,7 +254,10 @@ class BackupDumper {
 		$table = ( $history == WikiExporter::CURRENT ) ? 'page' : 'revision';
 		$field = ( $history == WikiExporter::CURRENT ) ? 'page_id' : 'rev_id';
 
-		$dbr = wfGetDB( DB_SLAVE );
+		$dbr = $this->forcedDb;
+		if ( $this->forcedDb === null ) {
+			$dbr = wfGetDB( DB_SLAVE );
+		}
 		$this->maxCount = $dbr->selectField( $table, "MAX($field)", '', __METHOD__ );
 		$this->startTime = wfTime();
 		$this->lastTime = $this->startTime;
@@ -259,6 +271,10 @@ class BackupDumper {
 	 * @return DatabaseBase
 	 */
 	function backupDb() {
+		if ( $this->forcedDb !== null ) {
+			return $this->forcedDb;
+		}
+
 		$this->lb = wfGetLBFactory()->newMainLB();
 		$db = $this->lb->getConnection( DB_SLAVE, 'backup' );
 
@@ -269,6 +285,18 @@ class BackupDumper {
 		return $db;
 	}
 
+	/**
+	 * Force the dump to use the provided database connection for database
+	 * operations, wherever possible.
+	 *
+	 * @param $db DatabaseBase|null: (Optional) the database connection to
+	 *            use. If null, resort to use the globally provided ways to
+	 *            get database connections.
+	 */
+	function setDb( DatabaseBase $db = null ) {
+		$this->forcedDb = $db;
+	}
+
 	function __destruct() {
 		if ( isset( $this->lb ) ) {
 			$this->lb->closeAll();
diff --git a/maintenance/backupTextPass.inc b/maintenance/backupTextPass.inc
index ecf2afb2ef..2b533ec4af 100644
--- a/maintenance/backupTextPass.inc
+++ b/maintenance/backupTextPass.inc
@@ -84,6 +84,11 @@ class TextPassDumper extends BackupDumper {
 			unset( $this->lb );
 		}
 
+		if ( $this->forcedDb !== null ) {
+			$this->db = $this->forcedDb;
+			return;
+		}
+
 		if ( isset( $this->db ) && $this->db->isOpen() ) {
 			throw new MWException( 'DB is set and has not been closed by the Load Balancer' );
 		}
-- 
2.20.1