Skip to main content

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