fix(postgres): prevent custom type stream deadlock

This commit is contained in:
Guoyu Su authored and GitHub committed 2026-09-03 19:28:20 +08:00
1 parent 0448d5bd2c
commit 2324d84c4e
2 files changed
+106

No files matched your search

+28
View File
@@ -1413,12 +1413,34 @@ fn postgres_select_stream_outcome(stream: tokio_postgres::RowStream) -> Postgres
PostgresSelectStreamOutcome::Binary { stream, metadata }
}
async fn prepare_unnamed_select_metadata(
client: &deadpool_postgres::Client,
sql: &str,
) -> Result<PreparedSelectMetadata, tokio_postgres::Error> {
// query_typed_raw sends Describe and Execute together. If its result has an
// unknown user-defined type, tokio-postgres then resolves that type with a
// second query on the same connection before returning the RowStream. A
// large result can fill the connection buffers and block that type lookup.
// Describe once without executing so custom types are cached before the
// actual unnamed stream starts. The stream remains unnamed to tolerate
// server-side prepared statement loss in snapshot/proxy environments.
let stmt = client.prepare(sql).await?;
Ok(prepared_select_metadata(stmt.columns()))
}
async fn start_postgres_select_stream(
client: &deadpool_postgres::Client,
sql: &str,
force_unnamed: bool,
) -> Result<PostgresSelectStreamOutcome, tokio_postgres::Error> {
if force_unnamed || postgres_client_uses_unnamed_statements(client) {
let metadata = prepare_unnamed_select_metadata(client, sql).await?;
if let Some(unsupported_type) = metadata.unsupported_type {
return Ok(PostgresSelectStreamOutcome::TextFallback {
column_types: metadata.column_types,
unsupported_type,
});
}
return postgres_query_unnamed(client, sql).await.map(postgres_select_stream_outcome);
}
@@ -13243,6 +13265,12 @@ mod tests {
.expect("stream unnamed query inside backup snapshot");
assert_eq!(columns, vec!["value"]);
assert_eq!(rows[0][0], serde_json::json!(45));
let prepared_count = client
.query_typed_one("SELECT count(*)::int8 FROM pg_prepared_statements", &[])
.await
.expect("count server-side prepared statements")
.get::<_, i64>(0);
assert_eq!(prepared_count, 0, "snapshot stream metadata preparation must not retain a named statement");
client.execute_typed("SELECT 1", &[]).await.expect("snapshot remains usable");
client.execute_typed("ROLLBACK", &[]).await.expect("rollback backup snapshot");
}