Skip to content
Snippets Groups Projects
Commit c963adca authored by Maksim Kaitmazian's avatar Maksim Kaitmazian
Browse files

fix(pgproto): correct row count for SELECT command tag

Previously, we retrieved the row count using rows.len(), where rows was an
iterator that returned the remaining number of rows instead of the
total number of rows sent.
parent bee0aa00
No related branches found
No related tags found
1 merge request!1155fix(pgproto): correct row count for SELECT command tag
Pipeline #45469 passed
......@@ -37,10 +37,6 @@ impl Rows {
RowDescription::new(self.desc.iter().map(Into::into).collect())
}
pub fn row_count(&self) -> usize {
self.rows.len()
}
pub fn values(&self) -> Vec<Vec<PgValue>> {
self.rows.clone().collect()
}
......@@ -66,5 +62,9 @@ pub enum ExecuteResult {
tag: CommandTag,
/// Rows we'll send to the client.
rows: Rows,
/// Cached number of rows in result.
/// Note: Rows is an iterator that contains only remaining rows. So it's
/// necessary to cache the number of rows before retrieving them.
row_count: usize,
},
}
......@@ -462,13 +462,15 @@ impl Portal {
}
}
PortalState::Running(ref mut stored_rows) => {
let taken = stored_rows.take(max_rows).collect();
let taken: Vec<_> = stored_rows.take(max_rows).collect();
let row_count = taken.len();
let rows = Rows::new(taken, self.describe.row_info());
if stored_rows.len() == 0 {
self.state = PortalState::Finished(None);
return Ok(ExecuteResult::FinishedDql {
rows,
tag: self.describe.command_tag(),
row_count,
});
}
return Ok(ExecuteResult::SuspendedDql { rows });
......
......@@ -50,13 +50,16 @@ pub fn process_execute_message(
&tag, row_count,
))?;
}
ExecuteResult::FinishedDql { tag, mut rows } => {
ExecuteResult::FinishedDql {
tag,
mut rows,
row_count,
} => {
while let Some(row) = rows.encode_next()? {
stream.write_message_noflush(messages::data_row(row))?;
}
stream.write_message_noflush(messages::command_complete_with_row_count(
&tag,
rows.row_count(),
&tag, row_count,
))?;
}
ExecuteResult::SuspendedDql { mut rows } => {
......
......@@ -17,15 +17,16 @@ pub fn process_query_message(
ExecuteResult::Dml { tag, row_count } => {
stream.write_message(messages::command_complete_with_row_count(&tag, row_count))?;
}
ExecuteResult::FinishedDql { tag, mut rows } => {
ExecuteResult::FinishedDql {
tag,
mut rows,
row_count,
} => {
stream.write_message_noflush(messages::row_description(rows.describe()))?;
while let Some(row) = rows.encode_next()? {
stream.write_message_noflush(messages::data_row(row))?;
}
stream.write_message(messages::command_complete_with_row_count(
&tag,
rows.row_count(),
))?;
stream.write_message(messages::command_complete_with_row_count(&tag, row_count))?;
}
ExecuteResult::SuspendedDql { .. } => {
return Err(PgError::InternalError(
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment