tank_postgres/
sql_writer.rs

1use std::{collections::BTreeMap, fmt::Write};
2use tank_core::{ColumnDef, Context, SqlWriter, Value, future::Either, separated_by};
3use time::{Date, OffsetDateTime, PrimitiveDateTime, Time};
4
5pub struct PostgresSqlWriter {}
6
7impl SqlWriter for PostgresSqlWriter {
8    fn as_dyn(&self) -> &dyn SqlWriter {
9        self
10    }
11
12    fn write_column_overridden_type(
13        &self,
14        _context: &mut Context,
15        out: &mut String,
16        _column: &ColumnDef,
17        types: &BTreeMap<&'static str, &'static str>,
18    ) {
19        if let Some(t) = types.iter().find_map(|(k, v)| {
20            if *k == "postgres" || *k == "postgresql" {
21                Some(v)
22            } else {
23                None
24            }
25        }) {
26            out.push_str(t);
27        }
28    }
29
30    fn write_column_type(&self, context: &mut Context, out: &mut String, value: &Value) {
31        match value {
32            Value::Boolean(..) => out.push_str("BOOLEAN"),
33            Value::Int8(..) => out.push_str("SMALLINT"),
34            Value::Int16(..) => out.push_str("SMALLINT"),
35            Value::Int32(..) => out.push_str("INTEGER"),
36            Value::Int64(..) => out.push_str("BIGINT"),
37            Value::Int128(..) => out.push_str("NUMERIC(39)"),
38            Value::UInt8(..) => out.push_str("SMALLINT"),
39            Value::UInt16(..) => out.push_str("INTEGER"),
40            Value::UInt32(..) => out.push_str("BIGINT"),
41            Value::UInt64(..) => out.push_str("NUMERIC(19)"),
42            Value::UInt128(..) => out.push_str("NUMERIC(39)"),
43            Value::Float32(..) => out.push_str("FLOAT4"),
44            Value::Float64(..) => out.push_str("FLOAT8"),
45            Value::Decimal(.., precision, scale) => {
46                out.push_str("NUMERIC");
47                if (precision, scale) != (&0, &0) {
48                    let _ = write!(out, "({},{})", precision, scale);
49                }
50            }
51            Value::Char(..) => out.push_str("CHAR(1)"),
52            Value::Varchar(..) => out.push_str("TEXT"),
53            Value::Blob(..) => out.push_str("BYTEA"),
54            Value::Date(..) => out.push_str("DATE"),
55            Value::Time(..) => out.push_str("TIME"),
56            Value::Timestamp(..) => out.push_str("TIMESTAMP"),
57            Value::TimestampWithTimezone(..) => out.push_str("TIMESTAMP WITH TIME ZONE"),
58            Value::Interval(..) => out.push_str("INTERVAL"),
59            Value::Uuid(..) => out.push_str("UUID"),
60            Value::Array(.., inner, size) => {
61                self.write_column_type(context, out, inner);
62                let _ = write!(out, "[{}]", size);
63            }
64            Value::List(.., inner) => {
65                self.write_column_type(context, out, inner);
66                out.push_str("[]");
67            }
68            _ => log::error!(
69                "Unexpected tank::Value, variant {:?} is not supported",
70                value
71            ),
72        };
73    }
74
75    fn write_value_blob(&self, _context: &mut Context, out: &mut String, value: &[u8]) {
76        out.push_str("'\\x");
77        for b in value {
78            let _ = write!(out, "{:X}", b);
79        }
80        out.push('\'');
81    }
82
83    fn write_value_date(
84        &self,
85        _context: &mut Context,
86        out: &mut String,
87        value: &Date,
88        timestamp: bool,
89    ) {
90        let (l, r) = if timestamp {
91            ("", "")
92        } else {
93            ("'", "'::DATE")
94        };
95        let (year, suffix) = if !timestamp && value.year() <= 0 {
96            // Year 0 in Postgres is 1 BC
97            (value.year().abs() + 1, " BC")
98        } else {
99            (value.year(), "")
100        };
101        let _ = write!(
102            out,
103            "{l}{:04}-{:02}-{:02}{suffix}{r}",
104            year,
105            value.month() as u8,
106            value.day()
107        );
108    }
109
110    fn write_value_time(
111        &self,
112        _context: &mut Context,
113        out: &mut String,
114        value: &Time,
115        timestamp: bool,
116    ) {
117        let mut subsecond = value.nanosecond();
118        let mut width = 9;
119        while width > 1 && subsecond % 10 == 0 {
120            subsecond /= 10;
121            width -= 1;
122        }
123        let (l, r) = if timestamp {
124            ("", "")
125        } else {
126            ("'", "'::TIME")
127        };
128        let _ = write!(
129            out,
130            "{l}{:02}:{:02}:{:02}.{:0width$}{r}",
131            value.hour(),
132            value.minute(),
133            value.second(),
134            subsecond
135        );
136    }
137
138    fn write_value_timestamp(
139        &self,
140        context: &mut Context,
141        out: &mut String,
142        value: &PrimitiveDateTime,
143    ) {
144        out.push('\'');
145        self.write_value_date(context, out, &value.date(), true);
146        out.push('T');
147        self.write_value_time(context, out, &value.time(), true);
148        if value.date().year() <= 0 {
149            out.push_str(" BC");
150        }
151        out.push_str("'::TIMESTAMP");
152    }
153
154    fn write_value_timestamptz(
155        &self,
156        context: &mut Context,
157        out: &mut String,
158        value: &OffsetDateTime,
159    ) {
160        out.push('\'');
161        self.write_value_date(context, out, &value.date(), true);
162        out.push('T');
163        self.write_value_time(context, out, &value.time(), true);
164        let _ = write!(
165            out,
166            "{:+03}:{:02}",
167            value.offset().whole_hours(),
168            value.offset().whole_minutes() % 60
169        );
170        if value.date().year() <= 0 {
171            out.push_str(" BC");
172        }
173        out.push_str("'::TIMESTAMPTZ");
174    }
175
176    fn write_value_list(
177        &self,
178        context: &mut Context,
179        out: &mut String,
180        value: Either<&Box<[Value]>, &Vec<Value>>,
181        ty: &Value,
182    ) {
183        out.push_str("ARRAY[");
184        separated_by(
185            out,
186            match value {
187                Either::Left(v) => v.iter(),
188                Either::Right(v) => v.iter(),
189            },
190            |out, v| {
191                self.write_value(context, out, v);
192            },
193            ",",
194        );
195        out.push_str("]::");
196        self.write_column_type(context, out, ty);
197    }
198
199    fn write_expression_operand_question_mark(&self, context: &mut Context, out: &mut String) {
200        context.counter += 1;
201        let _ = write!(out, "${}", context.counter);
202    }
203}