diff --git a/crates/dbx-core/src/sql.rs b/crates/dbx-core/src/sql.rs index c4c90f829d..63fa54108f 100644 --- a/crates/dbx-core/src/sql.rs +++ b/crates/dbx-core/src/sql.rs @@ -1369,6 +1369,20 @@ pub fn prepare_sql_file_statement( driver_profile: Option<&str>, ) -> SqlFileStatementAction { let statement = statement.trim(); + let postgres_statement = + (*db_type == DatabaseType::Postgres).then(|| strip_postgres_pg_dump_guard_command(statement)).flatten(); + let statement = postgres_statement.as_deref().unwrap_or(statement).trim(); + if statement.is_empty() + || (*db_type == DatabaseType::Postgres + && postgres_statement.is_some() + && !has_executable_sql_with_options( + statement, + SqlParsingOptions::for_database_type(DatabaseType::Postgres), + )) + { + return SqlFileStatementAction::Skip; + } + let is_mysql_compatible_target = is_mysql_compatible_import_target(db_type, driver_profile); if is_mysql_compatible_target && is_mysql_lock_table_statement(statement) { return SqlFileStatementAction::Skip; @@ -1393,6 +1407,37 @@ pub fn prepare_sql_file_statement( SqlFileStatementAction::Execute(body.to_string()) } +fn strip_postgres_pg_dump_guard_command(statement: &str) -> Option { + let executable = + leading_executable_sql_with_options(statement, SqlParsingOptions::for_database_type(DatabaseType::Postgres)); + let line_end = executable.find('\n').unwrap_or(executable.len()); + if !is_postgres_pg_dump_guard_command(&executable[..line_end]) { + return None; + } + + let prefix_len = statement.len() - executable.len(); + let suffix_start = prefix_len + line_end + usize::from(line_end < executable.len()); + let mut prepared = String::with_capacity(statement.len()); + prepared.push_str(&statement[..prefix_len]); + prepared.push_str(&statement[suffix_start..]); + Some(prepared) +} + +fn is_postgres_pg_dump_guard_command(line: &str) -> bool { + let mut parts = line.split_whitespace(); + let Some(command) = parts.next() else { + return false; + }; + if !matches!(command, "\\restrict" | "\\unrestrict") { + return false; + } + + let Some(key) = parts.next() else { + return false; + }; + parts.next().is_none() && key.bytes().all(|byte| byte.is_ascii_alphanumeric()) +} + pub fn optimize_sql_file_import_statements( statements: &[String], db_type: Option, @@ -2825,6 +2870,42 @@ mod tests { ); } + #[test] + fn prepares_postgres_pg_dump_restrict_guard_with_following_sql() { + assert_eq!( + prepare_sql_file_statement( + "-- PostgreSQL database dump\n\n\\restrict PreICrqH3VX69RwvK06tauVFPoZ3tThfV0y2zIFL2uSfhfNPSu1pWtu1QTudUAc\n\nSET statement_timeout = 0", + &DatabaseType::Postgres, + None + ), + SqlFileStatementAction::Execute("-- PostgreSQL database dump\n\n\nSET statement_timeout = 0".to_string()) + ); + } + + #[test] + fn skips_standalone_postgres_pg_dump_unrestrict_guard() { + assert_eq!( + prepare_sql_file_statement( + "--\n-- PostgreSQL database dump complete\n--\n\n\\unrestrict PreICrqH3VX69RwvK06tauVFPoZ3tThfV0y2zIFL2uSfhfNPSu1pWtu1QTudUAc", + &DatabaseType::Postgres, + None + ), + SqlFileStatementAction::Skip + ); + } + + #[test] + fn keeps_other_postgres_psql_commands_unsupported() { + assert_eq!( + prepare_sql_file_statement("\\copy users FROM 'users.csv'", &DatabaseType::Postgres, None), + SqlFileStatementAction::Execute("\\copy users FROM 'users.csv'".to_string()) + ); + assert_eq!( + prepare_sql_file_statement("\\restrict invalid-key", &DatabaseType::Postgres, None), + SqlFileStatementAction::Execute("\\restrict invalid-key".to_string()) + ); + } + #[test] fn optimizes_adjacent_single_row_inserts_into_multi_row_insert() { let statements = vec![ diff --git a/crates/dbx-core/src/sql_file_import.rs b/crates/dbx-core/src/sql_file_import.rs index 01f16b98b9..8a3b79fba4 100644 --- a/crates/dbx-core/src/sql_file_import.rs +++ b/crates/dbx-core/src/sql_file_import.rs @@ -1659,6 +1659,22 @@ mod tests { ); } + #[test] + fn postgres_pg_dump_guards_are_removed_from_import_plan() { + let statements = split_sql_file_import_statements( + "-- PostgreSQL database dump\n\n\\restrict Guard123\n\nSET statement_timeout = 0;\nSELECT 1;\n\n-- PostgreSQL database dump complete\n\n\\unrestrict Guard123", + Some(DatabaseType::Postgres), + ); + let planned = optimize_sql_file_import_statements(&statements, Some(DatabaseType::Postgres), None); + + assert_eq!(planned.len(), 3); + assert_eq!(planned[0].kind, SqlFileImportStatementKind::Execute); + assert!(!planned[0].sql.contains("\\restrict")); + assert!(planned[0].sql.contains("SET statement_timeout = 0")); + assert_eq!(planned[1].sql, "SELECT 1"); + assert_eq!(planned[2].kind, SqlFileImportStatementKind::Skip); + } + #[test] fn stop_on_error_returns_err_with_terminal_error_progress() { let decision = statement_error_decision(