servers/postgres/types/
bytea.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use bytes::BufMut;
16use pgwire::types::ToSqlText;
17use pgwire::types::format::FormatOptions;
18use postgres_types::{IsNull, ToSql, Type};
19
20#[derive(Debug)]
21pub struct HexOutputBytea<'a>(pub &'a [u8]);
22impl ToSqlText for HexOutputBytea<'_> {
23    fn to_sql_text(
24        &self,
25        ty: &Type,
26        out: &mut bytes::BytesMut,
27        format_options: &FormatOptions,
28    ) -> std::result::Result<IsNull, Box<dyn std::error::Error + Sync + Send>>
29    where
30        Self: Sized,
31    {
32        let _ = self.0.to_sql_text(ty, out, format_options);
33        Ok(IsNull::No)
34    }
35}
36
37impl ToSql for HexOutputBytea<'_> {
38    fn to_sql(
39        &self,
40        ty: &Type,
41        out: &mut bytes::BytesMut,
42    ) -> std::result::Result<IsNull, Box<dyn std::error::Error + Sync + Send>>
43    where
44        Self: Sized,
45    {
46        self.0.to_sql(ty, out)
47    }
48
49    fn accepts(ty: &Type) -> bool
50    where
51        Self: Sized,
52    {
53        <&[u8] as ToSql>::accepts(ty)
54    }
55
56    fn to_sql_checked(
57        &self,
58        ty: &Type,
59        out: &mut bytes::BytesMut,
60    ) -> std::result::Result<IsNull, Box<dyn std::error::Error + Sync + Send>> {
61        self.0.to_sql_checked(ty, out)
62    }
63}
64#[derive(Debug)]
65pub struct EscapeOutputBytea<'a>(pub &'a [u8]);
66impl ToSqlText for EscapeOutputBytea<'_> {
67    fn to_sql_text(
68        &self,
69        _ty: &Type,
70        out: &mut bytes::BytesMut,
71        _format_options: &FormatOptions,
72    ) -> std::result::Result<IsNull, Box<dyn std::error::Error + Sync + Send>>
73    where
74        Self: Sized,
75    {
76        self.0.iter().for_each(|b| match b {
77            0..=31 | 127..=255 => {
78                out.put_slice(b"\\");
79                out.put_slice(format!("{:03o}", b).as_bytes());
80            }
81            92 => out.put_slice(b"\\\\"),
82            32..=126 => out.put_u8(*b),
83        });
84        Ok(IsNull::No)
85    }
86}
87impl ToSql for EscapeOutputBytea<'_> {
88    fn to_sql(
89        &self,
90        ty: &Type,
91        out: &mut bytes::BytesMut,
92    ) -> std::result::Result<IsNull, Box<dyn std::error::Error + Sync + Send>>
93    where
94        Self: Sized,
95    {
96        self.0.to_sql(ty, out)
97    }
98
99    fn accepts(ty: &Type) -> bool
100    where
101        Self: Sized,
102    {
103        <&[u8] as ToSql>::accepts(ty)
104    }
105
106    fn to_sql_checked(
107        &self,
108        ty: &Type,
109        out: &mut bytes::BytesMut,
110    ) -> std::result::Result<IsNull, Box<dyn std::error::Error + Sync + Send>> {
111        self.0.to_sql_checked(ty, out)
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    #[test]
120    fn test_escape_output_bytea() {
121        let input: &[u8] = &[97, 98, 99, 107, 108, 109, 42, 169, 84];
122        let input = EscapeOutputBytea(input);
123
124        let expected = b"abcklm*\\251T";
125        let mut out = bytes::BytesMut::new();
126        let is_null = input
127            .to_sql_text(&Type::BYTEA, &mut out, &FormatOptions::default())
128            .unwrap();
129        assert!(matches!(is_null, IsNull::No));
130        assert_eq!(&out[..], expected);
131
132        let expected = &[97, 98, 99, 107, 108, 109, 42, 169, 84];
133        let mut out = bytes::BytesMut::new();
134        let is_null = input.to_sql(&Type::BYTEA, &mut out).unwrap();
135        assert!(matches!(is_null, IsNull::No));
136        assert_eq!(&out[..], expected);
137    }
138
139    #[test]
140    fn test_hex_output_bytea() {
141        let input = b"hello, world!";
142        let input = HexOutputBytea(input);
143
144        let expected = b"\\x68656c6c6f2c20776f726c6421";
145        let mut out = bytes::BytesMut::new();
146        let is_null = input
147            .to_sql_text(&Type::BYTEA, &mut out, &FormatOptions::default())
148            .unwrap();
149        assert!(matches!(is_null, IsNull::No));
150        assert_eq!(&out[..], expected);
151
152        let expected = b"hello, world!";
153        let mut out = bytes::BytesMut::new();
154        let is_null = input.to_sql(&Type::BYTEA, &mut out).unwrap();
155        assert!(matches!(is_null, IsNull::No));
156        assert_eq!(&out[..], expected);
157    }
158}
servers/postgres/types/bytea.rs

servers/postgres/types/
bytea.rs
