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
7pub struct PostgresSqlWriter {}
11
12impl PostgresSqlWriter {
13 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 (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}