tank_postgres/
sql_writer.rs

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

tank_postgres/
sql_writer.rs
