33use super :: { Migration , MigrationRecord , MigrationRunner } ;
44use crate :: Result ;
55use chrono:: Utc ;
6+ use sqlparser:: { dialect:: PostgreSqlDialect , parser:: Parser } ;
67use sqlx:: { PgPool , Row } ;
7- use tracing:: { debug, info} ;
8+ use tracing:: { debug, info, warn} ;
9+
10+ /// Parse SQL text into individual statements using sqlparser-rs
11+ /// This properly handles quoted strings, dollar-quoted strings, and comments
12+ fn parse_sql_statements ( sql : & str ) -> Result < Vec < String > , sqlparser:: parser:: ParserError > {
13+ let dialect = PostgreSqlDialect { } ;
14+
15+ // First, try to parse the entire SQL as a series of statements
16+ match Parser :: parse_sql ( & dialect, sql) {
17+ Ok ( statements) => {
18+ // Convert parsed statements back to SQL strings
19+ Ok ( statements. iter ( ) . map ( |stmt| format ! ( "{};" , stmt) ) . collect ( ) )
20+ }
21+ Err ( _) => {
22+ // If parsing fails, try to split manually while respecting SQL syntax
23+ // This is a more conservative approach for complex migrations
24+ Ok ( split_sql_respecting_quotes ( sql) )
25+ }
26+ }
27+ }
28+
29+ /// Split SQL while respecting quoted contexts
30+ /// This is a simpler fallback that handles the most common cases
31+ fn split_sql_respecting_quotes ( sql : & str ) -> Vec < String > {
32+ let mut statements = Vec :: new ( ) ;
33+ let mut current_statement = String :: new ( ) ;
34+ let mut in_single_quote = false ;
35+ let mut in_dollar_quote = false ;
36+ let mut dollar_tag = String :: new ( ) ;
37+ let mut chars = sql. chars ( ) . peekable ( ) ;
38+
39+ while let Some ( ch) = chars. next ( ) {
40+ current_statement. push ( ch) ;
41+
42+ match ch {
43+ '\'' if !in_dollar_quote => {
44+ // Handle single quotes (with escape sequences)
45+ if in_single_quote {
46+ // Check for escaped quote
47+ if chars. peek ( ) == Some ( & '\'' ) {
48+ current_statement. push ( chars. next ( ) . unwrap ( ) ) ;
49+ } else {
50+ in_single_quote = false ;
51+ }
52+ } else {
53+ in_single_quote = true ;
54+ }
55+ }
56+ '$' if !in_single_quote => {
57+ // Handle dollar quoting
58+ if in_dollar_quote {
59+ // Check if this closes the dollar quote
60+ let mut temp_tag = String :: new ( ) ;
61+ let chars_ahead: Vec < char > = chars. clone ( ) . collect ( ) ;
62+ let mut i = 0 ;
63+
64+ while i < chars_ahead. len ( ) && ( chars_ahead[ i] . is_alphanumeric ( ) || chars_ahead[ i] == '_' ) {
65+ temp_tag. push ( chars_ahead[ i] ) ;
66+ i += 1 ;
67+ }
68+
69+ if i < chars_ahead. len ( ) && chars_ahead[ i] == '$' && temp_tag == dollar_tag {
70+ // Consume the tag and closing $
71+ for _ in 0 ..=i {
72+ if let Some ( c) = chars. next ( ) {
73+ current_statement. push ( c) ;
74+ }
75+ }
76+ in_dollar_quote = false ;
77+ dollar_tag. clear ( ) ;
78+ }
79+ } else {
80+ // Check if this starts a dollar quote
81+ let mut temp_tag = String :: new ( ) ;
82+ let chars_ahead: Vec < char > = chars. clone ( ) . collect ( ) ;
83+ let mut i = 0 ;
84+
85+ while i < chars_ahead. len ( ) && ( chars_ahead[ i] . is_alphanumeric ( ) || chars_ahead[ i] == '_' ) {
86+ temp_tag. push ( chars_ahead[ i] ) ;
87+ i += 1 ;
88+ }
89+
90+ if i < chars_ahead. len ( ) && chars_ahead[ i] == '$' {
91+ // This is a dollar quote start
92+ for _ in 0 ..=i {
93+ if let Some ( c) = chars. next ( ) {
94+ current_statement. push ( c) ;
95+ }
96+ }
97+ in_dollar_quote = true ;
98+ dollar_tag = temp_tag;
99+ }
100+ }
101+ }
102+ ';' if !in_single_quote && !in_dollar_quote => {
103+ // This is a statement terminator
104+ let trimmed = current_statement. trim ( ) ;
105+ if !trimmed. is_empty ( ) && !trimmed. starts_with ( "--" ) {
106+ statements. push ( current_statement. clone ( ) ) ;
107+ }
108+ current_statement. clear ( ) ;
109+ }
110+ _ => {
111+ // Regular character, just continue
112+ }
113+ }
114+ }
115+
116+ // Add final statement if non-empty
117+ let trimmed = current_statement. trim ( ) ;
118+ if !trimmed. is_empty ( ) && !trimmed. starts_with ( "--" ) {
119+ statements. push ( current_statement) ;
120+ }
121+
122+ statements
123+ }
8124
9125/// PostgreSQL migration runner
10126pub struct PostgresMigrationRunner {
@@ -24,13 +140,18 @@ impl MigrationRunner<sqlx::Postgres> for PostgresMigrationRunner {
24140
25141 let mut tx = self . pool . begin ( ) . await ?;
26142
27- // Split SQL into individual statements and execute each one
28- // Split on semicolon and filter out empty statements
29- let statements: Vec < & str > = sql
30- . split ( ';' )
31- . map ( |s| s. trim ( ) )
32- . filter ( |s| !s. is_empty ( ) && !s. chars ( ) . all ( |c| c. is_whitespace ( ) || c == '\n' ) )
33- . collect ( ) ;
143+ // Parse SQL using sqlparser-rs for proper statement splitting
144+ let statements = match parse_sql_statements ( sql) {
145+ Ok ( stmts) => stmts,
146+ Err ( e) => {
147+ warn ! ( "Failed to parse SQL with sqlparser, falling back to naive splitting: {}" , e) ;
148+ // Fallback to simple splitting for compatibility
149+ sql. split ( ';' )
150+ . map ( |s| s. trim ( ) . to_string ( ) )
151+ . filter ( |s| !s. is_empty ( ) && !s. chars ( ) . all ( |c| c. is_whitespace ( ) || c == '\n' ) )
152+ . collect ( )
153+ }
154+ } ;
34155
35156 for ( i, statement) in statements. iter ( ) . enumerate ( ) {
36157 // Add semicolon back if it was removed by split
0 commit comments