Add support for MySQL user variables

Author: JanJakes

tests/WP_SQLite_Driver_Tests.php

@@ -6055,4 +6055,18 @@ public function testSetStatement(): void {
 		$this->assertQuery( 'SET CHARSET utf8mb4' );
 		$this->assertQuery( 'SET CHARACTER SET utf8mb4' );
 	}
+
+	public function testUserVariables(): void {
+		$this->assertQuery( 'SET @my_var = 1' );
+		$result = $this->assertQuery( 'SELECT @my_var' );
+		$this->assertEquals( 1, $result[0]->{'@my_var'} );
+
+		$this->assertQuery( 'SET @my_var = @my_var + 1' );
+		$result = $this->assertQuery( 'SELECT @my_var' );
+		$this->assertEquals( 2, $result[0]->{'@my_var'} );
+
+		$this->assertQuery( 'SET @my_var = @my_var + 1' );
+		$result = $this->assertQuery( 'SELECT @my_var' );
+		$this->assertEquals( 3, $result[0]->{'@my_var'} );
+	}
 }

wp-includes/sqlite-ast/class-wp-sqlite-driver.php

@@ -425,6 +425,18 @@ class WP_SQLite_Driver {
 		'STRICT_TRANS_TABLES',
 	);
 
+	/**
+	 * A name-to-value map of MySQL user variables.
+	 *
+	 * MySQL user variables are session-specific, so we can store them in-memory.
+	 *
+	 * See:
+	 *   https://dev.mysql.com/doc/refman/8.4/en/user-variables.html
+	 *
+	 * @var array<string, string>
+	 */
+	private $user_variables = array();
+
 	/**
 	 * Constructor.
 	 *
@@ -2109,11 +2121,19 @@ private function execute_set_statement( WP_Parser_Node $node ): void {
 					|| 'setSystemVariable' === $part->rule_name
 				)
 			) {
+				// Set a system variable.
 				array_shift( $definition ); // Remove the '='.
 				$value = array_shift( $definition );
 				$this->execute_set_system_variable_statement( $part, $value, $default_type );
+			} elseif (
+				$part instanceof WP_Parser_Node
+				&& 'userVariable' === $part->rule_name
+			) {
+				// Set a user variable.
+				array_shift( $definition ); // Remove the '='.
+				$value = array_shift( $definition );
+				$this->execute_set_user_variable_statement( $part, $value );
 			} else {
-				// TODO: Support user variables (in-memory or a temporary table).
 				throw $this->new_not_supported_exception(
 					sprintf( 'SET statement: %s', $node->rule_name )
 				);
@@ -2176,6 +2196,26 @@ private function execute_set_system_variable_statement(
 		// TODO: Handle GLOBAL, PERSIST, and PERSIST_ONLY types.
 	}
 
+	/**
+	 * Translate and execute a MySQL SET statement for user variables.
+	 *
+	 * @param  WP_Parser_Node $user_variable The "userVariable" AST node.
+	 * @param  WP_Parser_Node $expr          The "expr" AST node.
+	 * @throws WP_SQLite_Driver_Exception    When the query execution fails.
+	 */
+	private function execute_set_user_variable_statement(
+		WP_Parser_Node $user_variable,
+		WP_Parser_Node $expr
+	): void {
+		$name  = $this->unquote_sqlite_identifier(
+			$this->translate( $user_variable->get_first_child() )
+		);
+		$name  = strtolower( substr( $name, 1 ) ); // Remove '@', normalize case.
+		$value = $this->evaluate_expression( $expr );
+
+		$this->user_variables[ $name ] = $value;
+	}
+
 	/**
 	 * Translate and execute a MySQL administration statement in SQLite.
 	 *
@@ -2258,6 +2298,33 @@ private function execute_administration_statement( WP_Parser_Node $node ): void
 		$this->set_results_from_fetched_data( $results );
 	}
 
+	/**
+	 * Evaluate an expression and return the value, preserving its type.
+	 *
+	 * This is used to support expressions in SET statements for MySQL variables.
+	 *
+	 * @param  WP_Parser_Node $node The "expr" AST node.
+	 * @return mixed                The value of the expression.
+	 */
+	public function evaluate_expression( WP_Parser_Node $node ) {
+		// To support expressions, we'll use a SQLite query.
+		$stmt = $this->execute_sqlite_query(
+			sprintf( 'SELECT %s', $this->translate( $node ) )
+		);
+
+		// MySQL variables are typed, so we need to preserve the value type.
+		$value = $stmt->fetchColumn();
+		$type  = $stmt->getColumnMeta( 0 )['native_type'];
+		if ( 'null' === $type ) {
+			return null;
+		} elseif ( 'integer' === $type ) {
+			return (int) $value;
+		} elseif ( 'double' === $type ) {
+			return (float) $value;
+		}
+		return $value;
+	}
+
 	/**
 	 * Translate a MySQL AST node or token to an SQLite query fragment.
 	 *
@@ -2428,6 +2495,17 @@ private function translate( $node ): ?string {
 				//        See: https://dev.mysql.com/doc/refman/8.4/en/server-system-variable-reference.html
 				//        See: https://dev.mysql.com/doc/refman/8.4/en/server-system-variables.html
 				return $value;
+			case 'userVariable':
+				$name  = $this->unquote_sqlite_identifier( $this->translate( $node->get_first_child() ) );
+				$name  = strtolower( substr( $name, 1 ) ); // Remove '@', normalize case.
+				$value = $this->user_variables[ $name ] ?? null;
+				if ( null === $value ) {
+					return 'NULL';
+				}
+				if ( is_string( $value ) ) {
+					return $this->connection->quote( $value );
+				}
+				return (string) $value;
 			case 'castType':
 				// Translate "CAST(... AS BINARY)" to "CAST(... AS BLOB)".
 				if ( $node->has_child_token( WP_MySQL_Lexer::BINARY_SYMBOL ) ) {