tank_postgres/
sql_writer.rs

1use std::{collections::BTreeMap, fmt::Write};
2use tank_core::{
3    ColumnDef, Context, Dataset, DynQuery, Entity, Expression, Fragment, SqlWriter, Value,
4    separated_by,
5};
6use time::{Date, OffsetDateTime, PrimitiveDateTime, Time};
7
8/// Postgres SQL writer.
9pub struct PostgresSqlWriter {}
10
11impl PostgresSqlWriter {
12    /// Write COPY FROM STDIN BINARY.
13    pub fn write_copy<'b, E>(&self, out: &mut DynQuery)
14    where
15        Self: Sized,
16        E: Entity + 'b,
17    {
18        out.buffer().reserve(128);
19        out.push_str("COPY ");
20        let mut context = Context::new(Default::default(), E::qualified_columns());
21        self.write_table_ref(&mut context, out, E::table());
22        out.push_str(" (");
23        separated_by(
24            out,
25            E::columns().iter(),
26            |out, col| {
27                self.write_identifier(&mut context, out, col.name(), true);
28            },
29            ", ",
30        );
31        out.push_str(") FROM STDIN BINARY;");
32    }
33}
34
35impl SqlWriter for PostgresSqlWriter {
36    fn as_dyn(&self) -> &dyn SqlWriter {
37        self
38    }
39
40    fn write_column_overridden_type(
41        &self,
42        _context: &mut Context,
43        out: &mut DynQuery,
44        _column: &ColumnDef,
45        types: &BTreeMap<&'static str, &'static str>,
46    ) {
47        if let Some(t) = types.iter().find_map(|(k, v)| {
48            if *k == "postgres" || *k == "postgresql" {
49                Some(v)
50            } else {
51                None
52            }
53        }) {
54            out.push_str(t);
55        }
56    }
57
58    fn write_column_type(&self, context: &mut Context, out: &mut DynQuery, value: &Value) {
59        match value {
60            Value::Boolean(..) => out.push_str("BOOLEAN"),
61            Value::Int8(..) => out.push_str("SMALLINT"),
62            Value::Int16(..) => out.push_str("SMALLINT"),
63            Value::Int32(..) => out.push_str("INTEGER"),
64            Value::Int64(..) => out.push_str("BIGINT"),
65            Value::Int128(..) => out.push_str("NUMERIC(39)"),
66            Value::UInt8(..) => out.push_str("SMALLINT"),
67            Value::UInt16(..) => out.push_str("INTEGER"),
68            Value::UInt32(..) => out.push_str("BIGINT"),
69            Value::UInt64(..) => out.push_str("NUMERIC(19)"),
70            Value::UInt128(..) => out.push_str("NUMERIC(39)"),
71            Value::Float32(..) => out.push_str("FLOAT4"),
72            Value::Float64(..) => out.push_str("FLOAT8"),
73            Value::Decimal(.., precision, scale) => {
74                out.push_str("NUMERIC");
75                if (precision, scale) != (&0, &0) {
76                    let _ = write!(out, "({},{})", precision, scale);
77                }
78            }
79            Value::Char(..) => out.push_str("CHAR(1)"),
80            Value::Varchar(..) => out.push_str("TEXT"),
81            Value::Blob(..) => out.push_str("BYTEA"),
82            Value::Date(..) => out.push_str("DATE"),
83            Value::Time(..) => out.push_str("TIME"),
84            Value::Timestamp(..) => out.push_str("TIMESTAMP"),
85            Value::TimestampWithTimezone(..) => out.push_str("TIMESTAMP WITH TIME ZONE"),
86            Value::Interval(..) => out.push_str("INTERVAL"),
87            Value::Uuid(..) => out.push_str("UUID"),
88            Value::Array(.., inner, size) => {
89                self.write_column_type(context, out, inner);
90                let _ = write!(out, "[{}]", size);
91            }
92            Value::List(.., inner) => {
93                self.write_column_type(context, out, inner);
94                out.push_str("[]");
95            }
96            Value::Map(..) | Value::Json(..) | Value::Struct(..) => out.push_str("JSON"),
97            _ => log::error!("Unexpected tank::Value, Postgres does not support {value:?}"),
98        };
99    }
100
101    fn write_blob(&self, _context: &mut Context, out: &mut DynQuery, value: &[u8]) {
102        out.push_str("'\\x");
103        for b in value {
104            let _ = write!(out, "{:02X}", b);
105        }
106        out.push('\'');
107    }
108
109    fn write_date(&self, context: &mut Context, out: &mut DynQuery, value: &Date) {
110        let (l, r) = match context.fragment {
111            Fragment::None | Fragment::ParameterBinding | Fragment::Timestamp => ("", ""),
112            Fragment::Json | Fragment::JsonKey => ("\"", "\""),
113            _ => ("'", "'::DATE"),
114        };
115        let (year, suffix) = if value.year() <= 0 {
116            // Year 0 in Postgres is 1 BC
117            let suffix = if context.fragment == Fragment::Timestamp {
118                ""
119            } else {
120                " BC"
121            };
122            (value.year().abs() + 1, suffix)
123        } else {
124            (value.year(), "")
125        };
126        let month = value.month() as u8;
127        let day = value.day();
128        let _ = write!(out, "{l}{year:04}-{month:02}-{day:02}{suffix}{r}");
129    }
130
131    fn write_time(&self, context: &mut Context, out: &mut DynQuery, value: &Time) {
132        let (l, r) = match context.fragment {
133            Fragment::None | Fragment::ParameterBinding | Fragment::Timestamp => ("", ""),
134            Fragment::Json | Fragment::JsonKey => ("\"", "\""),
135            _ => ("'", "'::TIME"),
136        };
137        let (h, m, s, ns) = value.as_hms_nano();
138        let mut subsecond = ns;
139        let mut width = 9;
140        while width > 1 && subsecond % 10 == 0 {
141            subsecond /= 10;
142            width -= 1;
143        }
144        let _ = write!(out, "{l}{h:02}:{m:02}:{s:02}.{subsecond:0width$}{r}",);
145    }
146
147    fn write_timestamp(
148        &self,
149        context: &mut Context,
150        out: &mut DynQuery,
151        value: &PrimitiveDateTime,
152    ) {
153        let is_timestamp = context.fragment == Fragment::Timestamp;
154        let (l, r) = match context.fragment {
155            Fragment::None | Fragment::ParameterBinding | Fragment::Timestamp => ("", ""),
156            Fragment::Json | Fragment::JsonKey => ("\"", "\""),
157            _ => ("'", "'::TIMESTAMP"),
158        };
159        let mut context = context.switch_fragment(Fragment::Timestamp);
160        out.push_str(l);
161        self.write_date(&mut context.current, out, &value.date());
162        out.push('T');
163        self.write_time(&mut context.current, out, &value.time());
164        if !is_timestamp && value.date().year() <= 0 {
165            out.push_str(" BC");
166        }
167        out.push_str(r);
168    }
169
170    fn write_timestamptz(&self, context: &mut Context, out: &mut DynQuery, value: &OffsetDateTime) {
171        let (l, r) = match context.fragment {
172            Fragment::None | Fragment::ParameterBinding | Fragment::Timestamp => ("", ""),
173            Fragment::Json | Fragment::JsonKey => ("\"", "\""),
174            _ => ("'", "'::TIMESTAMPTZ"),
175        };
176        let mut context = context.switch_fragment(Fragment::Timestamp);
177        out.push_str(l);
178        self.write_timestamp(
179            &mut context.current,
180            out,
181            &PrimitiveDateTime::new(value.date(), value.time()),
182        );
183        let total_minutes = value.offset().whole_minutes();
184        let sign = if total_minutes >= 0 { '+' } else { '-' };
185        let _ = write!(
186            out,
187            "{}{:02}:{:02}",
188            sign,
189            (total_minutes.abs() / 60) as u8,
190            (total_minutes.abs() % 60) as u8
191        );
192        if value.date().year() <= 0 {
193            out.push_str(" BC");
194        }
195        out.push_str(r);
196    }
197
198    fn write_list(
199        &self,
200        context: &mut Context,
201        out: &mut DynQuery,
202        value: &mut dyn Iterator<Item = &dyn Expression>,
203        ty: Option<&Value>,
204        _is_array: bool,
205    ) {
206        out.push_str("ARRAY[");
207        separated_by(
208            out,
209            value,
210            |out, v| {
211                v.write_query(self, context, out);
212            },
213            ",",
214        );
215        out.push(']');
216        if let Some(ty) = ty {
217            out.push_str("::");
218            self.write_column_type(context, out, ty);
219            out.push_str("[]");
220        }
221    }
222
223    fn write_question_mark(&self, context: &mut Context, out: &mut DynQuery) {
224        context.counter += 1;
225        let _ = write!(out, "${}", context.counter);
226    }
227
228    fn write_current_timestamp_ms(&self, _context: &mut Context, out: &mut DynQuery) {
229        out.push_str("CAST(EXTRACT(EPOCH FROM NOW()) * 1000 AS BIGINT)");
230    }
231}
tank_postgres/sql_writer.rs

tank_postgres/
sql_writer.rs
