fix(postgres): prevent custom type stream deadlock
This commit is contained in:
1 parent
0448d5bd2c
commit
2324d84c4e
2 files changed
+106
No files matched your search
@@ -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");
|
||||
}
|
||||
|
||||
Reference in new issue
Block a user