Parse queries from "wp_get_db_schema()"

Author: JanJakes

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

@@ -655,7 +655,7 @@ public function query( string $query, $fetch_mode = PDO::FETCH_OBJ, ...$fetch_mo
 
 		try {
 			// Parse the MySQL query.
-			$ast = $this->parse_query( $query );
+			$ast = $this->parse_query( $query )->current();
 
 			// Handle transaction commands.
 
@@ -728,23 +728,24 @@ public function query( string $query, $fetch_mode = PDO::FETCH_OBJ, ...$fetch_mo
 	}
 
 	/**
-	 * Parse a MySQL query into an AST.
+	 * Parse a MySQL query into an array of ASTs.
 	 *
-	 * @param  string         $query      The MySQL query to parse.
-	 * @return WP_Parser_Node             The AST representing the parsed query.
+	 * @param  string           $query    The MySQL query to parse.
+	 * @return Generator<WP_Parser_Node>  A generator of ASTs representing the queries parsed from the input.
 	 * @throws WP_SQLite_Driver_Exception When the query parsing fails.
 	 */
-	public function parse_query( string $query ): WP_Parser_Node {
+	public function parse_query( string $query ): Generator {
 		$lexer  = new WP_MySQL_Lexer( $query );
 		$tokens = $lexer->remaining_tokens();
 
 		$parser = new WP_MySQL_Parser( self::$mysql_grammar, $tokens );
-		$ast    = $parser->parse();
-
-		if ( null === $ast ) {
-			throw $this->new_driver_exception( 'Failed to parse the MySQL query.' );
+		while ( $parser->next_ast() ) {
+			$ast = $parser->get_ast();
+			if ( null === $ast ) {
+				throw $this->new_driver_exception( 'Failed to parse the MySQL query.' );
+			}
+			yield $ast;
 		}
-		return $ast;
 	}
 
 	/**

wp-includes/sqlite-ast/class-wp-sqlite-information-schema-reconstructor.php

@@ -64,12 +64,16 @@ public function ensure_correct_information_schema(): void {
 				throw new Exception( 'The "wp_get_db_schema()" function was not defined.' );
 			}
 			$schema = wp_get_db_schema();
-			$parts  = preg_split( '/(CREATE\s+TABLE)/', $schema, -1, PREG_SPLIT_NO_EMPTY );
-			foreach ( $parts as $part ) {
-				$name               = $this->unquote_mysql_identifier(
-					preg_split( '/\s+/', $part, 2, PREG_SPLIT_NO_EMPTY )[0]
-				);
-				$wp_tables[ $name ] = 'CREATE TABLE' . $part;
+			foreach ( $this->driver->parse_query( $schema ) as $query ) {
+				$create_node = $query->get_first_descendant_node( 'createStatement' );
+				if ( $create_node && $create_node->has_child_node( 'createTable' ) ) {
+					$name_node = $create_node->get_first_descendant_node( 'tableName' );
+					$name      = $this->unquote_mysql_identifier(
+						substr( $schema, $name_node->get_start(), $name_node->get_length() )
+					);
+
+					$wp_tables[ $name ] = $create_node;
+				}
 			}
 		}
 
@@ -78,12 +82,12 @@ public function ensure_correct_information_schema(): void {
 			if ( ! in_array( $table, $information_schema_tables, true ) ) {
 				if ( isset( $wp_tables[ $table ] ) ) {
 					// WordPress core table (as returned by "wp_get_db_schema()").
-					$sql = $wp_tables[ $table ];
+					$ast = $wp_tables[ $table ];
 				} else {
 					// Other table (a WordPress plugin or unrelated to WordPress).
 					$sql = $this->generate_create_table_statement( $table );
+					$ast = $this->driver->parse_query( $sql )->current();
 				}
-				$ast = $this->driver->parse_query( $sql );
 				$this->information_schema_builder->record_create_table( $ast );
 			}
 		}
@@ -92,7 +96,7 @@ public function ensure_correct_information_schema(): void {
 		foreach ( $information_schema_tables as $table ) {
 			if ( ! in_array( $table, $tables, true ) ) {
 				$sql = sprintf( 'DROP %s', $this->quote_sqlite_identifier( $table ) );
-				$ast = $this->driver->parse_query( $sql );
+				$ast = $this->driver->parse_query( $sql )->current();
 				$this->information_schema_builder->record_drop_table( $ast );
 			}
 		}