tank_postgres/
sql_writer.rs

1use std::{collections::BTreeMap, fmt::Write};
2use tank_core::{
3    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 String)
15    where
16        Self: Sized,
17        E: Entity + 'b,
18    {
19        out.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 String,
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 String, 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!(
99                "Unexpected tank::Value, variant {:?} is not supported",
100                value
101            ),
102        };
103    }
104
105    fn write_value_blob(&self, _context: &mut Context, out: &mut String, value: &[u8]) {
106        out.push_str("'\\x");
107        for b in value {
108            let _ = write!(out, "{:02X}", b);
109        }
110        out.push('\'');
111    }
112
113    fn write_value_date(
114        &self,
115        _context: &mut Context,
116        out: &mut String,
117        value: &Date,
118        timestamp: bool,
119    ) {
120        let (l, r) = if timestamp {
121            ("", "")
122        } else {
123            ("'", "'::DATE")
124        };
125        let (year, suffix) = if !timestamp && value.year() <= 0 {
126            // Year 0 in Postgres is 1 BC
127            (value.year().abs() + 1, " BC")
128        } else {
129            (value.year(), "")
130        };
131        let _ = write!(
132            out,
133            "{l}{:04}-{:02}-{:02}{suffix}{r}",
134            year,
135            value.month() as u8,
136            value.day()
137        );
138    }
139
140    fn write_value_time(
141        &self,
142        _context: &mut Context,
143        out: &mut String,
144        value: &Time,
145        timestamp: bool,
146    ) {
147        let mut subsecond = value.nanosecond();
148        let mut width = 9;
149        while width > 1 && subsecond % 10 == 0 {
150            subsecond /= 10;
151            width -= 1;
152        }
153        let (l, r) = if timestamp {
154            ("", "")
155        } else {
156            ("'", "'::TIME")
157        };
158        let _ = write!(
159            out,
160            "{l}{:02}:{:02}:{:02}.{:0width$}{r}",
161            value.hour(),
162            value.minute(),
163            value.second(),
164            subsecond
165        );
166    }
167
168    fn write_value_timestamp(
169        &self,
170        context: &mut Context,
171        out: &mut String,
172        value: &PrimitiveDateTime,
173    ) {
174        out.push('\'');
175        self.write_value_date(context, out, &value.date(), true);
176        out.push('T');
177        self.write_value_time(context, out, &value.time(), true);
178        if value.date().year() <= 0 {
179            out.push_str(" BC");
180        }
181        out.push_str("'::TIMESTAMP");
182    }
183
184    fn write_value_timestamptz(
185        &self,
186        context: &mut Context,
187        out: &mut String,
188        value: &OffsetDateTime,
189    ) {
190        out.push('\'');
191        self.write_value_date(context, out, &value.date(), true);
192        out.push('T');
193        self.write_value_time(context, out, &value.time(), true);
194        let _ = write!(
195            out,
196            "{:+03}:{:02}",
197            value.offset().whole_hours(),
198            value.offset().whole_minutes() % 60
199        );
200        if value.date().year() <= 0 {
201            out.push_str(" BC");
202        }
203        out.push_str("'::TIMESTAMPTZ");
204    }
205
206    fn write_value_list(
207        &self,
208        context: &mut Context,
209        out: &mut String,
210        value: Either<&Box<[Value]>, &Vec<Value>>,
211        ty: &Value,
212        _elem_ty: &Value,
213    ) {
214        out.push_str("ARRAY[");
215        separated_by(
216            out,
217            match value {
218                Either::Left(v) => v.iter(),
219                Either::Right(v) => v.iter(),
220            },
221            |out, v| {
222                self.write_value(context, out, v);
223            },
224            ",",
225        );
226        out.push_str("]::");
227        self.write_column_type(context, out, ty);
228    }
229
230    fn write_expression_operand_question_mark(&self, context: &mut Context, out: &mut String) {
231        context.counter += 1;
232        let _ = write!(out, "${}", context.counter);
233    }
234}
tank_postgres/sql_writer.rs

tank_postgres/
sql_writer.rs
