use bytes::{Buf, BufMut, Bytes, BytesMut};
use crate::common::ByteStr;
pub trait UsizeExt {
fn to_u32(self) -> u32;
fn to_u16(self) -> u16;
}
pub trait StrExt {
fn nul_string_len(&self) -> u32;
}
pub trait BufMutExt {
fn put_nul_string(&mut self, string: &str);
}
pub trait BytesExt {
fn get_nul_bytestr(&mut self) -> Result<ByteStr, std::str::Utf8Error>;
}
pub trait BindParams: Buf {
fn size(&self) -> i32;
}
pub trait FmtExt {
fn lossy(&self) -> LossyFmt<'_>;
}
pub struct LossyFmt<'a>(pub &'a [u8]);
impl UsizeExt for usize {
fn to_u32(self) -> u32 {
self.try_into().expect("message size too large for protocol: {err}")
}
fn to_u16(self) -> u16 {
self.try_into().expect("message size too large for protocol: {err}")
}
}
impl StrExt for str {
fn nul_string_len(&self) -> u32 {
self.len().to_u32() + 1
}
}
impl<B: BufMut> BufMutExt for B {
fn put_nul_string(&mut self, string: &str) {
self.put(string.as_bytes());
self.put_u8(b'\0');
}
}
impl BytesExt for Bytes {
fn get_nul_bytestr(&mut self) -> Result<ByteStr, std::str::Utf8Error> {
let end = self
.iter()
.position(|e| matches!(e, b'\0'))
.expect("Postgres string did not nul terminated");
let me = self.split_to(end);
Buf::advance(self, 1); ByteStr::from_utf8(me)
}
}
impl BytesExt for BytesMut {
fn get_nul_bytestr(&mut self) -> Result<ByteStr, std::str::Utf8Error> {
let end = self
.iter()
.position(|e| matches!(e, b'\0'))
.expect("Postgres string did not nul terminated");
let me = self.split_to(end);
Buf::advance(self, 1); ByteStr::from_utf8(me.freeze())
}
}
impl FmtExt for [u8] {
fn lossy(&self) -> LossyFmt<'_> {
LossyFmt(self)
}
}
impl std::fmt::Display for LossyFmt<'_> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
for &b in self.0 {
if b.is_ascii_graphic() || b.is_ascii_whitespace() {
write!(f, "{}", b as char)?;
} else {
write!(f, "\\x{b:x}")?;
}
}
Ok(())
}
}
impl std::fmt::Debug for LossyFmt<'_> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "b\"{self}\"")
}
}