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            (value.year().abs() + 1, " BC")
118        } else {
119            (value.year(), "")
120        };
121        let month = value.month() as u8;
122        let day = value.day();
123        let _ = write!(out, "{l}{year:04}-{month:02}-{day:02}{suffix}{r}");
124    }
125
126    fn write_time(&self, context: &mut Context, out: &mut DynQuery, value: &Time) {
127        let (l, r) = match context.fragment {
128            Fragment::None | Fragment::ParameterBinding | Fragment::Timestamp => ("", ""),
129            Fragment::Json | Fragment::JsonKey => ("\"", "\""),
130            _ => ("'", "'::TIME"),
131        };
132        let (h, m, s, ns) = value.as_hms_nano();
133        let mut subsecond = ns;
134        let mut width = 9;
135        while width > 1 && subsecond % 10 == 0 {
136            subsecond /= 10;
137            width -= 1;
138        }
139        let _ = write!(out, "{l}{h:02}:{m:02}:{s:02}.{subsecond:0width$}{r}",);
140    }
141
142    fn write_timestamp(
143        &self,
144        context: &mut Context,
145        out: &mut DynQuery,
146        value: &PrimitiveDateTime,
147    ) {
148        let (l, r) = match context.fragment {
149            Fragment::None | Fragment::ParameterBinding | Fragment::Timestamp => ("", ""),
150            Fragment::Json | Fragment::JsonKey => ("\"", "\""),
151            _ => ("'", "'::TIMESTAMP"),
152        };
153        let mut context = context.switch_fragment(Fragment::Timestamp);
154        out.push_str(l);
155        self.write_date(&mut context.current, out, &value.date());
156        out.push('T');
157        self.write_time(&mut context.current, out, &value.time());
158        if value.date().year() <= 0 {
159            out.push_str(" BC");
160        }
161        out.push_str(r);
162    }
163
164    fn write_timestamptz(&self, context: &mut Context, out: &mut DynQuery, value: &OffsetDateTime) {
165        let (l, r) = match context.fragment {
166            Fragment::None | Fragment::ParameterBinding | Fragment::Timestamp => ("", ""),
167            Fragment::Json | Fragment::JsonKey => ("\"", "\""),
168            _ => ("'", "'::TIMESTAMPTZ"),
169        };
170        let mut context = context.switch_fragment(Fragment::Timestamp);
171        out.push_str(l);
172        self.write_timestamp(
173            &mut context.current,
174            out,
175            &PrimitiveDateTime::new(value.date(), value.time()),
176        );
177        let _ = write!(
178            out,
179            "{:+03}:{:02}",
180            value.offset().whole_hours(),
181            value.offset().whole_minutes() % 60
182        );
183        if value.date().year() <= 0 {
184            out.push_str(" BC");
185        }
186        out.push_str(r);
187    }
188
189    fn write_list(
190        &self,
191        context: &mut Context,
192        out: &mut DynQuery,
193        value: &mut dyn Iterator<Item = &dyn Expression>,
194        ty: Option<&Value>,
195        _is_array: bool,
196    ) {
197        out.push_str("ARRAY[");
198        separated_by(
199            out,
200            value,
201            |out, v| {
202                v.write_query(self, context, out);
203            },
204            ",",
205        );
206        out.push(']');
207        if let Some(ty) = ty {
208            out.push_str("::");
209            self.write_column_type(context, out, ty);
210            out.push_str("[]");
211        }
212    }
213
214    fn write_question_mark(&self, context: &mut Context, out: &mut DynQuery) {
215        context.counter += 1;
216        let _ = write!(out, "${}", context.counter);
217    }
218
219    fn write_current_timestamp_ms(&self, _context: &mut Context, out: &mut DynQuery) {
220        out.push_str("CAST(EXTRACT(EPOCH FROM NOW()) * 1000 AS BIGINT)");
221    }
222}
tank_postgres/sql_writer.rs

tank_postgres/
sql_writer.rs
