tank_postgres/
sql_writer.rs

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