tank_postgres/
sql_writer.rs

1use std::{collections::BTreeMap, fmt::Write};
2use tank_core::{
3    ColumnDef, Context, Dataset, DynQuery, Entity, Expression, Fragment, SqlCoreWriter,
4    SqlExpressionWriter, SqlFragmentWriter, SqlValueWriter, SqlWriter, Value, 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<E>(&self, out: &mut DynQuery)
14    where
15        Self: Sized,
16        E: Entity,
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 SqlCoreWriter 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
102impl SqlValueWriter for PostgresSqlWriter {
103    fn write_blob(&self, _context: &mut Context, out: &mut DynQuery, value: &[u8]) {
104        out.push_str("'\\x");
105        for b in value {
106            let _ = write!(out, "{:02X}", b);
107        }
108        out.push('\'');
109    }
110
111    fn write_date(&self, context: &mut Context, out: &mut DynQuery, value: &Date) {
112        let (l, r) = match context.fragment {
113            Fragment::None | Fragment::ParameterBinding | Fragment::Timestamp => ("", ""),
114            Fragment::Json | Fragment::JsonKey => ("\"", "\""),
115            _ => ("'", "'::DATE"),
116        };
117        let (year, suffix) = if value.year() <= 0 {
118            // Year 0 in Postgres is 1 BC
119            let suffix = if context.fragment == Fragment::Timestamp {
120                ""
121            } else {
122                " BC"
123            };
124            (value.year().abs() + 1, suffix)
125        } else {
126            (value.year(), "")
127        };
128        let month = value.month() as u8;
129        let day = value.day();
130        let _ = write!(out, "{l}{year:04}-{month:02}-{day:02}{suffix}{r}");
131    }
132
133    fn write_time(&self, context: &mut Context, out: &mut DynQuery, value: &Time) {
134        let (l, r) = match context.fragment {
135            Fragment::None | Fragment::ParameterBinding | Fragment::Timestamp => ("", ""),
136            Fragment::Json | Fragment::JsonKey => ("\"", "\""),
137            _ => ("'", "'::TIME"),
138        };
139        let (h, m, s, ns) = value.as_hms_nano();
140        let mut subsecond = ns;
141        let mut width = 9;
142        while width > 1 && subsecond % 10 == 0 {
143            subsecond /= 10;
144            width -= 1;
145        }
146        let _ = write!(out, "{l}{h:02}:{m:02}:{s:02}.{subsecond:0width$}{r}",);
147    }
148
149    fn write_timestamp(
150        &self,
151        context: &mut Context,
152        out: &mut DynQuery,
153        value: &PrimitiveDateTime,
154    ) {
155        let is_timestamp = context.fragment == Fragment::Timestamp;
156        let (l, r) = match context.fragment {
157            Fragment::None | Fragment::ParameterBinding | Fragment::Timestamp => ("", ""),
158            Fragment::Json | Fragment::JsonKey => ("\"", "\""),
159            _ => ("'", "'::TIMESTAMP"),
160        };
161        let mut context = context.switch_fragment(Fragment::Timestamp);
162        out.push_str(l);
163        self.write_date(&mut context.current, out, &value.date());
164        out.push('T');
165        self.write_time(&mut context.current, out, &value.time());
166        if !is_timestamp && value.date().year() <= 0 {
167            out.push_str(" BC");
168        }
169        out.push_str(r);
170    }
171
172    fn write_timestamptz(&self, context: &mut Context, out: &mut DynQuery, value: &OffsetDateTime) {
173        let (l, r) = match context.fragment {
174            Fragment::None | Fragment::ParameterBinding | Fragment::Timestamp => ("", ""),
175            Fragment::Json | Fragment::JsonKey => ("\"", "\""),
176            _ => ("'", "'::TIMESTAMPTZ"),
177        };
178        let mut context = context.switch_fragment(Fragment::Timestamp);
179        out.push_str(l);
180        self.write_timestamp(
181            &mut context.current,
182            out,
183            &PrimitiveDateTime::new(value.date(), value.time()),
184        );
185        let (h, m, s) = value.offset().as_hms();
186        let sign = if h < 0 || m < 0 || s < 0 { '-' } else { '+' };
187        let _ = write!(out, "{sign}{:02}:{:02}", h.unsigned_abs(), m.unsigned_abs());
188        if s != 0 {
189            let _ = write!(out, ":{:02}", s.unsigned_abs());
190        }
191        if value.date().year() <= 0 {
192            out.push_str(" BC");
193        }
194        out.push_str(r);
195    }
196
197    fn write_list(
198        &self,
199        context: &mut Context,
200        out: &mut DynQuery,
201        value: &mut dyn Iterator<Item = &dyn Expression>,
202        ty: Option<&Value>,
203        _is_array: bool,
204    ) {
205        out.push_str("ARRAY[");
206        separated_by(
207            out,
208            value,
209            |out, v| {
210                v.write_query(self, context, out);
211            },
212            ",",
213        );
214        out.push(']');
215        if let Some(ty) = ty {
216            out.push_str("::");
217            self.write_column_type(context, out, ty);
218            out.push_str("[]");
219        }
220    }
221}
222
223impl SqlExpressionWriter for PostgresSqlWriter {
224    fn write_question_mark(&self, context: &mut Context, out: &mut DynQuery) {
225        context.counter += 1;
226        let _ = write!(out, "${}", context.counter);
227    }
228
229    fn write_current_timestamp_ms(&self, _context: &mut Context, out: &mut DynQuery) {
230        out.push_str("CAST(EXTRACT(EPOCH FROM NOW()) * 1000 AS BIGINT)");
231    }
232}
233
234impl SqlFragmentWriter for PostgresSqlWriter {}
235impl SqlWriter for PostgresSqlWriter {}
tank_postgres/sql_writer.rs

tank_postgres/
sql_writer.rs
