use bytes::{BufMut, BytesMut};
use crate::oid::TypeOid;
pub trait Encode {
fn type_oid(&self) -> TypeOid;
fn encode(&self, buf: &mut BytesMut);
fn encode_param(&self, buf: &mut BytesMut) {
let start = buf.len();
buf.put_i32(0); self.encode(buf);
let len = (buf.len() - start - 4) as i32;
let start_bytes = &mut buf[start..start + 4];
start_bytes.copy_from_slice(&len.to_be_bytes());
}
}
impl Encode for bool {
fn type_oid(&self) -> TypeOid {
TypeOid::Bool
}
fn encode(&self, buf: &mut BytesMut) {
buf.put_u8(if *self { 1 } else { 0 });
}
}
impl Encode for i16 {
fn type_oid(&self) -> TypeOid {
TypeOid::Int2
}
fn encode(&self, buf: &mut BytesMut) {
buf.put_i16(*self);
}
}
impl Encode for i32 {
fn type_oid(&self) -> TypeOid {
TypeOid::Int4
}
fn encode(&self, buf: &mut BytesMut) {
buf.put_i32(*self);
}
}
impl Encode for i64 {
fn type_oid(&self) -> TypeOid {
TypeOid::Int8
}
fn encode(&self, buf: &mut BytesMut) {
buf.put_i64(*self);
}
}
impl Encode for f32 {
fn type_oid(&self) -> TypeOid {
TypeOid::Float4
}
fn encode(&self, buf: &mut BytesMut) {
buf.put_f32(*self);
}
}
impl Encode for f64 {
fn type_oid(&self) -> TypeOid {
TypeOid::Float8
}
fn encode(&self, buf: &mut BytesMut) {
buf.put_f64(*self);
}
}
impl Encode for str {
fn type_oid(&self) -> TypeOid {
TypeOid::Text
}
fn encode(&self, buf: &mut BytesMut) {
buf.put_slice(self.as_bytes());
}
}
impl Encode for String {
fn type_oid(&self) -> TypeOid {
TypeOid::Text
}
fn encode(&self, buf: &mut BytesMut) {
buf.put_slice(self.as_bytes());
}
}
impl Encode for &str {
fn type_oid(&self) -> TypeOid {
TypeOid::Text
}
fn encode(&self, buf: &mut BytesMut) {
buf.put_slice(self.as_bytes());
}
}
impl Encode for [u8] {
fn type_oid(&self) -> TypeOid {
TypeOid::Bytea
}
fn encode(&self, buf: &mut BytesMut) {
buf.put_slice(self);
}
}
impl Encode for Vec<u8> {
fn type_oid(&self) -> TypeOid {
TypeOid::Bytea
}
fn encode(&self, buf: &mut BytesMut) {
buf.put_slice(self);
}
}
impl Encode for crate::newtypes::PgNumeric {
fn type_oid(&self) -> TypeOid {
TypeOid::Numeric
}
fn encode(&self, buf: &mut BytesMut) {
buf.put_slice(self.0.as_bytes());
}
}
impl Encode for crate::newtypes::PgTimestamp {
fn type_oid(&self) -> TypeOid {
TypeOid::Timestamp
}
fn encode(&self, buf: &mut BytesMut) {
match self {
crate::newtypes::PgTimestamp::Value(us) => buf.put_i64(*us),
crate::newtypes::PgTimestamp::Infinity => buf.put_i64(i64::MAX),
crate::newtypes::PgTimestamp::NegInfinity => buf.put_i64(i64::MIN),
}
}
}
impl Encode for crate::newtypes::PgDate {
fn type_oid(&self) -> TypeOid {
TypeOid::Date
}
fn encode(&self, buf: &mut BytesMut) {
match self {
crate::newtypes::PgDate::Value(d) => buf.put_i32(*d),
crate::newtypes::PgDate::Infinity => buf.put_i32(i32::MAX),
crate::newtypes::PgDate::NegInfinity => buf.put_i32(i32::MIN),
}
}
}
impl Encode for crate::newtypes::PgInet {
fn type_oid(&self) -> TypeOid {
TypeOid::Inet
}
fn encode(&self, buf: &mut BytesMut) {
buf.put_slice(self.0.as_bytes());
}
}
pub fn encode_array_header(buf: &mut BytesMut, has_null: bool, element_oid: u32, len: usize) {
buf.put_i32(1); buf.put_i32(if has_null { 1 } else { 0 });
buf.put_u32(element_oid);
buf.put_i32(len as i32); buf.put_i32(1); }
macro_rules! impl_array_encode {
($t:ty, $array_oid_variant:ident) => {
impl Encode for Vec<$t> {
fn type_oid(&self) -> TypeOid {
TypeOid::$array_oid_variant
}
fn encode(&self, buf: &mut BytesMut) {
encode_array_header(buf, false, <$t as crate::pg_type::PgType>::OID, self.len());
for v in self {
v.encode_param(buf);
}
}
}
};
}
impl_array_encode!(bool, BoolArray);
impl_array_encode!(i16, Int2Array);
impl_array_encode!(i32, Int4Array);
impl_array_encode!(i64, Int8Array);
impl_array_encode!(f32, Float4Array);
impl_array_encode!(f64, Float8Array);
impl_array_encode!(String, TextArray);
#[cfg(feature = "chrono")]
const PG_EPOCH_OFFSET_US: i64 = 946_684_800_000_000;
#[cfg(feature = "chrono")]
impl Encode for chrono::NaiveDateTime {
fn type_oid(&self) -> TypeOid {
TypeOid::Timestamp
}
fn encode(&self, buf: &mut BytesMut) {
let us = self.and_utc().timestamp_micros() - PG_EPOCH_OFFSET_US;
buf.put_i64(us);
}
}
#[cfg(feature = "chrono")]
impl Encode for chrono::DateTime<chrono::Utc> {
fn type_oid(&self) -> TypeOid {
TypeOid::Timestamptz
}
fn encode(&self, buf: &mut BytesMut) {
let us = self.timestamp_micros() - PG_EPOCH_OFFSET_US;
buf.put_i64(us);
}
}
#[cfg(feature = "chrono")]
impl Encode for chrono::NaiveDate {
fn type_oid(&self) -> TypeOid {
TypeOid::Date
}
fn encode(&self, buf: &mut BytesMut) {
let pg_epoch = chrono::NaiveDate::from_ymd_opt(2000, 1, 1).unwrap();
let days = (*self - pg_epoch).num_days() as i32;
buf.put_i32(days);
}
}
#[cfg(feature = "chrono")]
impl Encode for chrono::NaiveTime {
fn type_oid(&self) -> TypeOid {
TypeOid::Time
}
fn encode(&self, buf: &mut BytesMut) {
let us = self
.signed_duration_since(chrono::NaiveTime::from_hms_opt(0, 0, 0).unwrap())
.num_microseconds()
.unwrap_or(0);
buf.put_i64(us);
}
}
#[cfg(feature = "json")]
impl Encode for serde_json::Value {
fn type_oid(&self) -> TypeOid {
TypeOid::Jsonb
}
fn encode(&self, buf: &mut BytesMut) {
buf.put_u8(1);
let json_text =
serde_json::to_string(self).expect("serde_json::Value serialization cannot fail");
buf.put_slice(json_text.as_bytes());
}
}
#[cfg(feature = "uuid")]
impl Encode for uuid::Uuid {
fn type_oid(&self) -> TypeOid {
TypeOid::Uuid
}
fn encode(&self, buf: &mut BytesMut) {
buf.put_slice(self.as_bytes());
}
}
impl_array_encode!(crate::newtypes::PgNumeric, NumericArray);
impl_array_encode!(crate::newtypes::PgInet, InetArray);
#[cfg(feature = "chrono")]
impl_array_encode!(chrono::NaiveDate, DateArray);
#[cfg(feature = "chrono")]
impl_array_encode!(chrono::NaiveTime, TimeArray);
#[cfg(feature = "chrono")]
impl_array_encode!(chrono::NaiveDateTime, TimestampArray);
#[cfg(feature = "chrono")]
impl_array_encode!(chrono::DateTime<chrono::Utc>, TimestamptzArray);
#[cfg(feature = "uuid")]
impl_array_encode!(uuid::Uuid, UuidArray);
#[cfg(feature = "json")]
impl_array_encode!(serde_json::Value, JsonbArray);
pub trait SqlParam: Sync {
fn param_oid(&self) -> u32;
fn encode_param_value(&self) -> Option<BytesMut>;
}
impl<T: Encode + Sync> SqlParam for T {
fn param_oid(&self) -> u32 {
self.type_oid().as_u32()
}
fn encode_param_value(&self) -> Option<BytesMut> {
let mut buf = BytesMut::new();
self.encode(&mut buf);
Some(buf)
}
}
impl<T: Encode + Sync> SqlParam for Option<T> {
fn param_oid(&self) -> u32 {
match self {
Some(v) => v.type_oid().as_u32(),
None => 0,
}
}
fn encode_param_value(&self) -> Option<BytesMut> {
match self {
Some(v) => {
let mut buf = BytesMut::new();
v.encode(&mut buf);
Some(buf)
}
None => None,
}
}
}