From: Brad Jorsch <bjorsch@wikimedia.org>
Date: Thu, 8 Feb 2018 19:16:29 +0000 (-0500)
Subject: Database: Add batching to non-native insertSelect()
X-Git-Tag: 1.31.0-rc.0~483^2
X-Git-Url: http://git.cyclocoop.org/%7B%24admin_url%7Dcompta/comptes/journal.php?a=commitdiff_plain;h=aefb143a8e8c963d6fedb471bbdace284be88196;p=lhc%2Fweb%2Fwiklou.git

Database: Add batching to non-native insertSelect()

It would be easy for a call to nonNativeInsertSelect() to generate an
INSERT that's too big for the database to actually process. Add batching
to try to avoid that.

Bug: T160993
Change-Id: I1de994208d95926f0d75c0d7cab7b5fe1dd565c3
---

diff --git a/includes/libs/rdbms/database/Database.php b/includes/libs/rdbms/database/Database.php
index 572a798c06..fd7bfe96bd 100644
--- a/includes/libs/rdbms/database/Database.php
+++ b/includes/libs/rdbms/database/Database.php
@@ -236,6 +236,9 @@ abstract class Database implements IDatabase, IMaintainableDatabase, LoggerAware
 	/** @var TransactionProfiler */
 	protected $trxProfiler;
 
+	/** @var int */
+	protected $nonNativeInsertSelectBatchSize = 10000;
+
 	/**
 	 * Constructor and database handle and attempt to connect to the DB server
 	 *
@@ -278,6 +281,10 @@ abstract class Database implements IDatabase, IMaintainableDatabase, LoggerAware
 		$this->queryLogger = $params['queryLogger'];
 		$this->errorLogger = $params['errorLogger'];
 
+		if ( isset( $params['nonNativeInsertSelectBatchSize'] ) ) {
+			$this->nonNativeInsertSelectBatchSize = $params['nonNativeInsertSelectBatchSize'];
+		}
+
 		// Set initial dummy domain until open() sets the final DB/prefix
 		$this->currentDomain = DatabaseDomain::newUnspecified();
 
@@ -331,6 +338,7 @@ abstract class Database implements IDatabase, IMaintainableDatabase, LoggerAware
 	 *   - cliMode: Whether to consider the execution context that of a CLI script.
 	 *   - agent: Optional name used to identify the end-user in query profiling/logging.
 	 *   - srvCache: Optional BagOStuff instance to an APC-style cache.
+	 *   - nonNativeInsertSelectBatchSize: Optional batch size for non-native INSERT SELECT emulation.
 	 * @return Database|null If the database driver or extension cannot be found
 	 * @throws InvalidArgumentException If the database driver or extension cannot be found
 	 * @since 1.18
@@ -2504,12 +2512,43 @@ abstract class Database implements IDatabase, IMaintainableDatabase, LoggerAware
 			return false;
 		}
 
-		$rows = [];
-		foreach ( $res as $row ) {
-			$rows[] = (array)$row;
+		try {
+			$affectedRowCount = 0;
+			$this->startAtomic( $fname );
+			$rows = [];
+			$ok = true;
+			foreach ( $res as $row ) {
+				$rows[] = (array)$row;
+
+				// Avoid inserts that are too huge
+				if ( count( $rows ) >= $this->nonNativeInsertSelectBatchSize ) {
+					$ok = $this->insert( $destTable, $rows, $fname, $insertOptions );
+					if ( !$ok ) {
+						break;
+					}
+					$affectedRowCount += $this->affectedRows();
+					$rows = [];
+				}
+			}
+			if ( $rows && $ok ) {
+				$ok = $this->insert( $destTable, $rows, $fname, $insertOptions );
+				if ( $ok ) {
+					$affectedRowCount += $this->affectedRows();
+				}
+			}
+			if ( $ok ) {
+				$this->endAtomic( $fname );
+				$this->affectedRowCount = $affectedRowCount;
+			} else {
+				$this->rollback( $fname, self::FLUSHING_INTERNAL );
+				$this->affectedRowCount = 0;
+			}
+			return $ok;
+		} catch ( Exception $e ) {
+			$this->rollback( $fname, self::FLUSHING_INTERNAL );
+			$this->affectedRowCount = 0;
+			throw $e;
 		}
-
-		return $this->insert( $destTable, $rows, $fname, $insertOptions );
 	}
 
 	/**
diff --git a/tests/phpunit/includes/libs/rdbms/database/DatabaseSQLTest.php b/tests/phpunit/includes/libs/rdbms/database/DatabaseSQLTest.php
index ebf6e45199..888654411a 100644
--- a/tests/phpunit/includes/libs/rdbms/database/DatabaseSQLTest.php
+++ b/tests/phpunit/includes/libs/rdbms/database/DatabaseSQLTest.php
@@ -457,7 +457,7 @@ class DatabaseSQLTest extends PHPUnit\Framework\TestCase {
 			isset( $sql['selectOptions'] ) ? $sql['selectOptions'] : [],
 			isset( $sql['selectJoinConds'] ) ? $sql['selectJoinConds'] : []
 		);
-		$this->assertLastSqlDb( implode( '; ', [ $sqlSelect, $sqlInsert ] ), $dbWeb );
+		$this->assertLastSqlDb( implode( '; ', [ $sqlSelect, 'BEGIN', $sqlInsert, 'COMMIT' ] ), $dbWeb );
 	}
 
 	public static function provideInsertSelect() {
@@ -537,6 +537,30 @@ class DatabaseSQLTest extends PHPUnit\Framework\TestCase {
 		];
 	}
 
+	public function testInsertSelectBatching() {
+		$dbWeb = new DatabaseTestHelper( __CLASS__, [ 'cliMode' => false ] );
+		$rows = [];
+		for ( $i = 0; $i <= 25000; $i++ ) {
+			$rows[] = [ 'field' => $i ];
+		}
+		$dbWeb->forceNextResult( $rows );
+		$dbWeb->insertSelect(
+			'insert_table',
+			'select_table',
+			[ 'field' => 'field2' ],
+			'*',
+			__METHOD__
+		);
+		$this->assertLastSqlDb( implode( '; ', [
+			'SELECT field2 AS field FROM select_table WHERE *   FOR UPDATE',
+			'BEGIN',
+			"INSERT INTO insert_table (field) VALUES ('" . implode( "'),('", range( 0, 9999 ) ) . "')",
+			"INSERT INTO insert_table (field) VALUES ('" . implode( "'),('", range( 10000, 19999 ) ) . "')",
+			"INSERT INTO insert_table (field) VALUES ('" . implode( "'),('", range( 20000, 25000 ) ) . "')",
+			'COMMIT'
+		] ), $dbWeb );
+	}
+
 	/**
 	 * @dataProvider provideReplace
 	 * @covers Wikimedia\Rdbms\Database::replace