Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/sql/postgres/protocol/ArrayList.rs
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,9 @@ impl<'a> WriterContext for ArrayListCtx<'a> {
arr[i..i + bytes.len()].copy_from_slice(bytes);
Ok(())
}
fn truncate(mut self, offset: usize) {
self.array_mut().truncate(offset);
}
}

pub type Writer<'a> = NewWriter<ArrayListCtx<'a>>;
20 changes: 20 additions & 0 deletions src/sql/postgres/protocol/NewWriter.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,10 @@ pub trait WriterContext: Copy {
fn offset(self) -> usize;
fn write(self, bytes: &[u8]) -> Result<(), AnyPostgresError>;
fn pwrite(self, bytes: &[u8], offset: usize) -> Result<(), AnyPostgresError>;
/// Discard everything written at or after `offset` (a value previously
/// returned by [`offset`]). Used to roll back a partially-written message
/// when encoding fails partway through.
fn truncate(self, offset: usize);
}

#[derive(Copy, Clone)]
Expand Down Expand Up @@ -60,6 +64,22 @@ impl<C: WriterContext> NewWriter<C> {
C::pwrite(self.wrapped, data, i)
}

/// Run `f`; on error, discard every byte it wrote so a half-serialised
/// message (e.g. a Bind whose parameter coercion threw in JS) is never
/// left in the buffer to desync the next query's flush.
#[inline]
pub fn atomically<T>(
self,
f: impl FnOnce(Self) -> Result<T, AnyPostgresError>,
) -> Result<T, AnyPostgresError> {
let start = self.offset();
let result = f(self);
if result.is_err() {
C::truncate(self.wrapped, start);
}
result
}

pub fn int4(self, value: PostgresInt32) -> Result<(), AnyPostgresError> {
self.write(&value.to_be_bytes())
}
Expand Down
194 changes: 100 additions & 94 deletions src/sql_jsc/postgres/PostgresRequest.rs
Original file line number Diff line number Diff line change
Expand Up @@ -289,66 +289,70 @@ pub(crate) fn prepare_and_query_with_signature<Context: WriterContext>(
global: &JSGlobalObject,
query: &[u8],
array_value: JSValue,
mut writer: protocol::NewWriter<Context>,
writer: protocol::NewWriter<Context>,
signature: &mut Signature,
) -> Result<(), AnyPostgresError> {
write_query(
query,
&signature.prepared_statement_name,
&signature.fields,
writer,
)?;
write_bind(
&signature.prepared_statement_name,
BunString::empty(),
global,
array_value,
JSValue::ZERO,
&[],
&[],
writer,
)?;
let exec = protocol::Execute {
p: protocol::PortalOrPreparedStatement::PreparedStatement(
writer.atomically(|mut writer| {
write_query(
query,
&signature.prepared_statement_name,
),
..Default::default()
};
exec.write_internal(&mut writer)?;
&signature.fields,
writer,
)?;
write_bind(
&signature.prepared_statement_name,
BunString::empty(),
global,
array_value,
JSValue::ZERO,
&[],
&[],
writer,
)?;
let exec = protocol::Execute {
p: protocol::PortalOrPreparedStatement::PreparedStatement(
&signature.prepared_statement_name,
),
..Default::default()
};
exec.write_internal(&mut writer)?;

writer.write(&protocol::FLUSH)?;
writer.write(&protocol::SYNC)?;
Ok(())
writer.write(&protocol::FLUSH)?;
writer.write(&protocol::SYNC)?;
Ok(())
})
}

pub(crate) fn bind_and_execute<Context: WriterContext>(
global: &JSGlobalObject,
statement: &PostgresSQLStatement,
array_value: JSValue,
columns_value: JSValue,
mut writer: protocol::NewWriter<Context>,
writer: protocol::NewWriter<Context>,
) -> Result<(), AnyPostgresError> {
write_bind(
&statement.signature.prepared_statement_name,
BunString::empty(),
global,
array_value,
columns_value,
&statement.parameters,
&statement.fields,
writer,
)?;
let exec = protocol::Execute {
p: protocol::PortalOrPreparedStatement::PreparedStatement(
writer.atomically(|mut writer| {
write_bind(
&statement.signature.prepared_statement_name,
),
..Default::default()
};
exec.write_internal(&mut writer)?;
BunString::empty(),
global,
array_value,
columns_value,
&statement.parameters,
&statement.fields,
writer,
)?;
let exec = protocol::Execute {
p: protocol::PortalOrPreparedStatement::PreparedStatement(
&statement.signature.prepared_statement_name,
),
..Default::default()
};
exec.write_internal(&mut writer)?;

writer.write(&protocol::FLUSH)?;
writer.write(&protocol::SYNC)?;
Ok(())
writer.write(&protocol::FLUSH)?;
writer.write(&protocol::SYNC)?;
Ok(())
})
}

/// Atomically sends Parse + [Describe] + Bind + Execute + Flush + Sync as a single message batch.
Expand All @@ -363,61 +367,63 @@ pub fn parse_and_bind_and_execute<Context: WriterContext>(
array_value: JSValue,
columns_value: JSValue,
include_describe: bool,
mut writer: protocol::NewWriter<Context>,
writer: protocol::NewWriter<Context>,
) -> Result<(), AnyPostgresError> {
let name = &statement.signature.prepared_statement_name;
writer.atomically(|mut writer| {
let name = &statement.signature.prepared_statement_name;

// Parse
{
let q = protocol::Parse {
name,
params: &statement.signature.fields,
query,
// Parse
{
let q = protocol::Parse {
name,
params: &statement.signature.fields,
query,
};
q.write_internal(&mut writer)?;
bun_core::scoped_log!(Postgres, "Parse: {}", bun_fmt::quote(query));
}

// Describe (needed on first execution to learn parameter/result types for caching)
if include_describe {
let d = protocol::Describe {
p: protocol::PortalOrPreparedStatement::PreparedStatement(name),
};
d.write_internal(writer)?;
bun_core::scoped_log!(Postgres, "Describe: {}", bun_fmt::quote(name));
}

// Bind — use server-provided types if available (binary format), otherwise
// fall back to signature types (text format for unknowns). The server will
// handle text-to-type conversion based on the parameter types from Parse.
let param_fields = if !statement.parameters.is_empty() {
&statement.parameters[..]
} else {
&statement.signature.fields[..]
};
q.write_internal(&mut writer)?;
bun_core::scoped_log!(Postgres, "Parse: {}", bun_fmt::quote(query));
}
let result_fields = &statement.fields;

// Describe (needed on first execution to learn parameter/result types for caching)
if include_describe {
let d = protocol::Describe {
write_bind(
name,
BunString::empty(),
global,
array_value,
columns_value,
param_fields,
result_fields,
writer,
)?;

// Execute
let exec = protocol::Execute {
p: protocol::PortalOrPreparedStatement::PreparedStatement(name),
..Default::default()
};
d.write_internal(writer)?;
bun_core::scoped_log!(Postgres, "Describe: {}", bun_fmt::quote(name));
}
exec.write_internal(&mut writer)?;

// Bind — use server-provided types if available (binary format), otherwise
// fall back to signature types (text format for unknowns). The server will
// handle text-to-type conversion based on the parameter types from Parse.
let param_fields = if !statement.parameters.is_empty() {
&statement.parameters[..]
} else {
&statement.signature.fields[..]
};
let result_fields = &statement.fields;

write_bind(
name,
BunString::empty(),
global,
array_value,
columns_value,
param_fields,
result_fields,
writer,
)?;

// Execute
let exec = protocol::Execute {
p: protocol::PortalOrPreparedStatement::PreparedStatement(name),
..Default::default()
};
exec.write_internal(&mut writer)?;

writer.write(&protocol::FLUSH)?;
writer.write(&protocol::SYNC)?;
Ok(())
writer.write(&protocol::FLUSH)?;
writer.write(&protocol::SYNC)?;
Ok(())
})
}

pub(crate) fn execute_query<Context: WriterContext>(
Expand Down
7 changes: 7 additions & 0 deletions src/sql_jsc/postgres/PostgresSQLConnection.rs
Original file line number Diff line number Diff line change
Expand Up @@ -1685,6 +1685,13 @@ impl protocol::WriterContext for Writer {
fn pwrite(mut self, bytes: &[u8], i: usize) -> Result<(), AnyPostgresError> {
Writer::pwrite(&mut self, bytes, i)
}
#[inline]
fn truncate(self, offset: usize) {
self.connection.write_buffer.with_mut(|b| {
debug_assert!(b.head as usize + offset <= b.byte_list.len());
b.byte_list.truncate(b.head as usize + offset);
});
}
}

impl PostgresSQLConnection {
Expand Down
Loading
Loading