Skip to content

Commit

Permalink
Merge pull request #1101 from Lantern-chat/async_size
Browse files Browse the repository at this point in the history
Shrink query_opt/query_one codegen size very slightly
  • Loading branch information
sfackler authored Mar 17, 2024
2 parents e528e01 + 9d7c43c commit 64caf4c
Showing 1 changed file with 16 additions and 20 deletions.
36 changes: 16 additions & 20 deletions tokio-postgres/src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -274,19 +274,9 @@ impl Client {
where
T: ?Sized + ToStatement,
{
let stream = self.query_raw(statement, slice_iter(params)).await?;
pin_mut!(stream);

let row = match stream.try_next().await? {
Some(row) => row,
None => return Err(Error::row_count()),
};

if stream.try_next().await?.is_some() {
return Err(Error::row_count());
}

Ok(row)
self.query_opt(statement, params)
.await
.and_then(|res| res.ok_or_else(Error::row_count))
}

/// Executes a statements which returns zero or one rows, returning it.
Expand All @@ -310,16 +300,22 @@ impl Client {
let stream = self.query_raw(statement, slice_iter(params)).await?;
pin_mut!(stream);

let row = match stream.try_next().await? {
Some(row) => row,
None => return Ok(None),
};
let mut first = None;

// Originally this was two calls to `try_next().await?`,
// once for the first element, and second to error if more than one.
//
// However, this new form with only one .await in a loop generates
// slightly smaller codegen/stack usage for the resulting future.
while let Some(row) = stream.try_next().await? {
if first.is_some() {
return Err(Error::row_count());
}

if stream.try_next().await?.is_some() {
return Err(Error::row_count());
first = Some(row);
}

Ok(Some(row))
Ok(first)
}

/// The maximally flexible version of [`query`].
Expand Down

0 comments on commit 64caf4c

Please sign in to comment.