From e498fc373120c28c69ddb02a461f150c90d0f34c Mon Sep 17 00:00:00 2001 From: zipg Date: Tue, 15 Sep 2026 23:22:23 +0800 Subject: [PATCH] fix(postgres): recover applied create table after disconnect --- crates/dbx-core/src/db/postgres.rs | 2 +- crates/dbx-core/src/query.rs | 201 +++++++++++++++++++++++++++++ 2 files changed, 202 insertions(+), 1 deletion(-) diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index fd6c228af6..df07aba681 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -7168,7 +7168,7 @@ pub async fn get_columns(pool: &Pool, schema: &str, table: &str) -> Result String { +pub(crate) fn pg_quote_literal(value: &str) -> String { format!("'{}'", value.replace('\'', "''")) } diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index a03d88acb5..a4d015bd6f 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -1508,6 +1508,54 @@ fn query_pool_error_action(db_type: Option, sql: &str, err: &str) } } +fn native_postgres_compatibility_type(db_type: Option) -> bool { + matches!( + db_type, + Some( + DatabaseType::Postgres + | DatabaseType::Gaussdb + | DatabaseType::OpenGauss + | DatabaseType::Kwdb + | DatabaseType::Questdb + ) + ) +} + +fn postgres_create_table_relation(sql: &str) -> Option<(Option, String)> { + let statements = Parser::parse_sql(&PostgreSqlDialect {}, sql).ok()?; + let [Statement::CreateTable(table)] = statements.as_slice() else { + return None; + }; + if table.temporary { + return None; + } + // PostgreSQL stores unquoted identifiers lower-cased, so the existence + // probe must compare the folded spelling; quoted names keep their case. + let mut parts = table.name.0.iter().filter_map(|part| { + part.as_ident().map(|ident| match ident.quote_style { + Some(_) => ident.value.clone(), + None => ident.value.to_lowercase(), + }) + }); + let table_name = parts.next_back()?; + Some((parts.next_back(), table_name)) +} + +fn should_verify_postgres_create_table_after_connection_error( + db_type: Option, + sql: &str, + error: &str, +) -> bool { + native_postgres_compatibility_type(db_type) + && is_connection_error(error) + && postgres_create_table_relation(sql).is_some() +} + +fn is_postgres_duplicate_relation_error(error: &str) -> bool { + let lower = error.to_ascii_lowercase(); + lower.contains("already exists") && (lower.contains("relation") || lower.contains("table")) +} + fn query_execution_error_action( db_type: Option, sql: &str, @@ -2839,6 +2887,85 @@ pub async fn execute_sql_statement_with_options_typed( result } +async fn recover_postgres_create_table_after_connection_error( + state: &AppState, + connection_id: &str, + database: &str, + schema: Option<&str>, + sql: &str, + cancel_token: Option, + options: &QueryExecutionOptions, + db_type: Option, + initial_error: &QueryExecutionError, +) -> Option> { + if is_canceled(&cancel_token) + || !should_verify_postgres_create_table_after_connection_error(db_type, sql, &initial_error.to_string()) + { + return None; + } + + let pool_database = query_pool_database(database, options.catalog.as_deref()); + let new_key = state + .reconnect_pool_for_session(connection_id, pool_database, options.client_session_id.as_deref()) + .await + .ok()?; + let mysql_dialect = connection_mysql_query_dialect(state, connection_id).await; + + // A compatible PostgreSQL server may close the session after committing a + // CREATE TABLE. Retry once on a fresh session: this covers a disconnect + // before execution, while a duplicate relation below confirms that the + // original request already took effect. + let retry_error = match do_execute_typed( + state, + &new_key, + mysql_dialect, + Some(database), + sql, + schema, + cancel_token.clone(), + options.clone(), + ) + .await + { + Ok(result) => return Some(Ok(result)), + Err(error) => error, + }; + if !is_postgres_duplicate_relation_error(&retry_error.to_string()) { + return None; + } + + let (qualified_schema, table_name) = postgres_create_table_relation(sql)?; + let schema_predicate = qualified_schema + .or_else(|| schema.map(str::to_owned)) + .map(|schema_name| format!("n.nspname = {}", db::postgres::pg_quote_literal(&schema_name))) + .unwrap_or_else(|| "n.nspname = current_schema()".to_string()); + let verify_sql = format!( + "SELECT EXISTS (SELECT 1 FROM pg_catalog.pg_class c JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace WHERE {} AND c.relname = {}) AS dbx_create_table_applied", + schema_predicate, + db::postgres::pg_quote_literal(&table_name), + ); + let verify_options = QueryExecutionOptions { + max_rows: Some(1), + client_session_id: options.client_session_id.clone(), + ..Default::default() + }; + let verified = do_execute_typed( + state, + &new_key, + mysql_dialect, + Some(database), + &verify_sql, + None, + cancel_token, + verify_options, + ) + .await + .ok() + .and_then(|result| result.rows.first().and_then(|row| row.first()).and_then(|value| value.as_bool())) + .unwrap_or(false); + verified.then(|| Ok(empty_query_result(0))) +} + async fn execute_sql_statement_with_options_typed_inner( state: &AppState, connection_id: &str, @@ -2903,6 +3030,23 @@ async fn execute_sql_statement_with_options_typed_inner( }; let action = result.as_ref().err().map(|error| query_execution_error_action(db_type, sql, error)); + if let Some(initial_error) = result.as_ref().err() { + if let Some(recovered) = recover_postgres_create_table_after_connection_error( + state, + connection_id, + database, + schema, + sql, + cancel_token.clone(), + &options, + db_type, + initial_error, + ) + .await + { + return with_sql_context(recovered); + } + } match action { Some(PoolErrorAction::ReconnectAndRetry) if !is_canceled(&cancel_token) => { let pool_database = query_pool_database(database, options.catalog.as_deref()); @@ -9620,6 +9764,63 @@ for line in sys.stdin: assert!(is_connection_error("Error occurred while creating a new object: error communicating with the server")); } + #[test] + fn postgres_create_table_recovery_only_targets_persistent_tables() { + assert_eq!( + postgres_create_table_relation("CREATE TABLE \"app\".\"events\" (id bigint)"), + Some((Some("app".to_string()), "events".to_string())) + ); + assert_eq!( + postgres_create_table_relation("-- ddl\nCREATE TABLE events (id bigint)"), + Some((None, "events".to_string())) + ); + assert_eq!(postgres_create_table_relation("CREATE TEMP TABLE events (id bigint)"), None); + assert_eq!(postgres_create_table_relation("ALTER TABLE events ADD COLUMN note text"), None); + } + + #[test] + fn postgres_create_table_recovery_folds_unquoted_identifiers() { + assert_eq!( + postgres_create_table_relation("CREATE TABLE MyTable (id bigint)"), + Some((None, "mytable".to_string())) + ); + assert_eq!( + postgres_create_table_relation("CREATE TABLE Core.Orders (id bigint)"), + Some((Some("core".to_string()), "orders".to_string())) + ); + assert_eq!( + postgres_create_table_relation("CREATE TABLE \"Core\".\"Orders\" (id bigint)"), + Some((Some("Core".to_string()), "Orders".to_string())) + ); + } + + #[test] + fn postgres_create_table_recovery_requires_native_connection_loss() { + let sql = "CREATE TABLE events (id bigint)"; + assert!(should_verify_postgres_create_table_after_connection_error( + Some(DatabaseType::Postgres), + sql, + "connection closed; PostgreSQL schema.reset cleanup failed: connection closed" + )); + assert!(should_verify_postgres_create_table_after_connection_error( + Some(DatabaseType::Gaussdb), + sql, + "connection reset by peer" + )); + assert!(!should_verify_postgres_create_table_after_connection_error( + Some(DatabaseType::Postgres), + sql, + "ERROR: permission denied" + )); + assert!(!should_verify_postgres_create_table_after_connection_error( + Some(DatabaseType::Vastbase), + sql, + "connection closed" + )); + assert!(is_postgres_duplicate_relation_error("ERROR: relation \"events\" already exists")); + assert!(!is_postgres_duplicate_relation_error("ERROR: permission denied for schema public")); + } + #[test] fn query_error_context_omits_raw_sql_and_is_not_duplicated() { let sql = "select 'secret-123' as token";