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
5pub struct PostgresSqlWriter {}
9
10impl PostgresSqlWriter {
11 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 (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}