Merge "rdbms: add query timeout support to Database::select()"

3 files changed, 57 insertions, 1 deletions

diff --git a/includes/libs/rdbms/database/DatabaseMysqlBase.php b/includes/libs/rdbms/database/DatabaseMysqlBase.php
index 2ae4dfe7c9f2..190e28878534 100644
--- a/includes/libs/rdbms/database/DatabaseMysqlBase.php
+++ b/includes/libs/rdbms/database/DatabaseMysqlBase.php
@@ -1380,7 +1380,36 @@ abstract class DatabaseMysqlBase extends Database {
 		return 'CAST( ' . $field . ' AS SIGNED )';
 	}
 
-	/**
+	public function selectSQLText(
+		$table,
+		$vars,
+		$conds = '',
+		$fname = __METHOD__,
+		$options = [],
+		$join_conds = []
+	) {
+		$sql = parent::selectSQLText( $table, $vars, $conds, $fname, $options, $join_conds );
+		// https://dev.mysql.com/doc/refman/5.7/en/optimizer-hints.html
+		// https://mariadb.com/kb/en/library/aborting-statements/
+		$timeoutMsec = intval( $options['MAX_EXECUTION_TIME'] ?? 0 );
+		if ( $timeoutMsec > 0 ) {
+			list( $vendor, $number ) = $this->getMySqlServerVariant();
+			if ( $vendor === 'MariaDB' && version_compare( $number, '10.1.2', '>=' ) ) {
+				$timeoutSec = $timeoutMsec / 1000;
+				$sql = "SET STATEMENT max_statement_time=$timeoutSec FOR $sql";
+			} elseif ( $vendor === 'MySQL' && version_compare( $number, '5.7.0', '>=' ) ) {
+				$sql = preg_replace(
+					'/^SELECT(?=\s)/',
+					"SELECT /*+ MAX_EXECUTION_TIME($timeoutMsec)*/",
+					$sql
+				);
+			}
+		}
+
+		return $sql;
+	}
+
+	/*
 	 * @return bool Whether GTID support is used (mockable for testing)
 	 */
 	protected function useGTIDs() {
diff --git a/includes/libs/rdbms/database/IDatabase.php b/includes/libs/rdbms/database/IDatabase.php
index e19dd7b8b27b..94ce6d8b290b 100644
--- a/includes/libs/rdbms/database/IDatabase.php
+++ b/includes/libs/rdbms/database/IDatabase.php
@@ -767,6 +767,9 @@ interface IDatabase {
 	 *   - EXPLAIN: In MySQL, this causes an EXPLAIN SELECT query to be run,
 	 *     instead of SELECT.
 	 *
+	 *   - MAX_EXECUTION_TIME: (only in MySQL/MariaDB) maximum allowed time to
+	 *     run the query in milliseconds (if database supports it).
+	 *
 	 * And also the following boolean MySQL extensions, see the MySQL manual
 	 * for documentation:
 	 *
diff --git a/tests/phpunit/unit/includes/libs/rdbms/database/DatabaseMysqlBaseTest.php b/tests/phpunit/unit/includes/libs/rdbms/database/DatabaseMysqlBaseTest.php
index 4fdd4ab732c2..1493056d2790 100644
--- a/tests/phpunit/unit/includes/libs/rdbms/database/DatabaseMysqlBaseTest.php
+++ b/tests/phpunit/unit/includes/libs/rdbms/database/DatabaseMysqlBaseTest.php
@@ -745,6 +745,30 @@ class DatabaseMysqlBaseTest extends PHPUnit\Framework\TestCase {
 	}
 
 	/**
+	 * @covers \Wikimedia\Rdbms\DatabaseMysqlBase::selectSQLText
+	 */
+	public function testMaxExecutionTime() {
+		$db = $this->getMockBuilder( DatabaseMysqli::class )
+			->disableOriginalConstructor()
+			->onlyMethods( [ 'getMySqlServerVariant', 'dbSchema', 'tablePrefix' ] )
+			->getMock();
+		$db->method( 'getMySqlServerVariant' )->willReturn( [ 'MariaDB', '10.4.21' ] );
+
+		/** @var IDatabase $db */
+		$sql = $db->selectSQLText( 'image',
+			'img_metadata',
+			'*',
+			'',
+			[ 'MAX_EXECUTION_TIME' => 1 ]
+		);
+
+		$this->assertEquals(
+			"SET STATEMENT max_statement_time=0.001 FOR SELECT  img_metadata  FROM `image`     ",
+			$sql
+		);
+	}
+
+	/**
 	 * @covers \Wikimedia\Rdbms\Database::streamStatementEnd
 	 * @covers \Wikimedia\Rdbms\DatabaseMysqlBase::streamStatementEnd
 	 */