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 postgres_types::{IsNull, ToSql, Type};
18
19#[derive(Debug)]
20pub struct HexOutputBytea<'a>(pub &'a [u8]);
21impl ToSqlText for HexOutputBytea<'_> {
22    fn to_sql_text(
23        &self,
24        ty: &Type,
25        out: &mut bytes::BytesMut,
26    ) -> std::result::Result<IsNull, Box<dyn std::error::Error + Sync + Send>>
27    where
28        Self: Sized,
29    {
30        let _ = self.0.to_sql_text(ty, out);
31        Ok(IsNull::No)
32    }
33}
34
35impl ToSql for HexOutputBytea<'_> {
36    fn to_sql(
37        &self,
38        ty: &Type,
39        out: &mut bytes::BytesMut,
40    ) -> std::result::Result<IsNull, Box<dyn std::error::Error + Sync + Send>>
41    where
42        Self: Sized,
43    {
44        self.0.to_sql(ty, out)
45    }
46
47    fn accepts(ty: &Type) -> bool
48    where
49        Self: Sized,
50    {
51        <&[u8] as ToSql>::accepts(ty)
52    }
53
54    fn to_sql_checked(
55        &self,
56        ty: &Type,
57        out: &mut bytes::BytesMut,
58    ) -> std::result::Result<IsNull, Box<dyn std::error::Error + Sync + Send>> {
59        self.0.to_sql_checked(ty, out)
60    }
61}
62#[derive(Debug)]
63pub struct EscapeOutputBytea<'a>(pub &'a [u8]);
64impl ToSqlText for EscapeOutputBytea<'_> {
65    fn to_sql_text(
66        &self,
67        _ty: &Type,
68        out: &mut bytes::BytesMut,
69    ) -> std::result::Result<IsNull, Box<dyn std::error::Error + Sync + Send>>
70    where
71        Self: Sized,
72    {
73        self.0.iter().for_each(|b| match b {
74            0..=31 | 127..=255 => {
75                out.put_slice(b"\\");
76                out.put_slice(format!("{:03o}", b).as_bytes());
77            }
78            92 => out.put_slice(b"\\\\"),
79            32..=126 => out.put_u8(*b),
80        });
81        Ok(IsNull::No)
82    }
83}
84impl ToSql for EscapeOutputBytea<'_> {
85    fn to_sql(
86        &self,
87        ty: &Type,
88        out: &mut bytes::BytesMut,
89    ) -> std::result::Result<IsNull, Box<dyn std::error::Error + Sync + Send>>
90    where
91        Self: Sized,
92    {
93        self.0.to_sql(ty, out)
94    }
95
96    fn accepts(ty: &Type) -> bool
97    where
98        Self: Sized,
99    {
100        <&[u8] as ToSql>::accepts(ty)
101    }
102
103    fn to_sql_checked(
104        &self,
105        ty: &Type,
106        out: &mut bytes::BytesMut,
107    ) -> std::result::Result<IsNull, Box<dyn std::error::Error + Sync + Send>> {
108        self.0.to_sql_checked(ty, out)
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115
116    #[test]
117    fn test_escape_output_bytea() {
118        let input: &[u8] = &[97, 98, 99, 107, 108, 109, 42, 169, 84];
119        let input = EscapeOutputBytea(input);
120
121        let expected = b"abcklm*\\251T";
122        let mut out = bytes::BytesMut::new();
123        let is_null = input.to_sql_text(&Type::BYTEA, &mut out).unwrap();
124        assert!(matches!(is_null, IsNull::No));
125        assert_eq!(&out[..], expected);
126
127        let expected = &[97, 98, 99, 107, 108, 109, 42, 169, 84];
128        let mut out = bytes::BytesMut::new();
129        let is_null = input.to_sql(&Type::BYTEA, &mut out).unwrap();
130        assert!(matches!(is_null, IsNull::No));
131        assert_eq!(&out[..], expected);
132    }
133
134    #[test]
135    fn test_hex_output_bytea() {
136        let input = b"hello, world!";
137        let input = HexOutputBytea(input);
138
139        let expected = b"\\x68656c6c6f2c20776f726c6421";
140        let mut out = bytes::BytesMut::new();
141        let is_null = input.to_sql_text(&Type::BYTEA, &mut out).unwrap();
142        assert!(matches!(is_null, IsNull::No));
143        assert_eq!(&out[..], expected);
144
145        let expected = b"hello, world!";
146        let mut out = bytes::BytesMut::new();
147        let is_null = input.to_sql(&Type::BYTEA, &mut out).unwrap();
148        assert!(matches!(is_null, IsNull::No));
149        assert_eq!(&out[..], expected);
150    }
151}
servers/postgres/types/bytea.rs

servers/postgres/types/
bytea.rs
