1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
#![allow(
clippy::big_endian_bytes,
reason = "PostgreSQL binary COPY format specifies big-endian encoding throughout"
)]
use chrono::{DateTime, Utc};
use futures::future::try_join_all;
use uuid::Uuid;
use crate::postgres::retry::{RetryPolicy, retry};
use crate::types::error::Result;
use crate::types::model::DataRow;
use super::PostgresRepository;
use sqlx::PgPool;
impl PostgresRepository {
/// Run a single insert pass via up to `workers` concurrent binary COPY streams:
/// the rows are split into `workers` chunks, each `COPYed` on its own
/// connection/backend. The streams are driven concurrently on one task, so while
/// one is blocked on socket I/O the others make progress. Returns the total
/// inserted.
///
/// Each partition is retried independently under a reduced-budget [`RetryPolicy`] so a
/// transient failure on one connection re-COPYs only that partition rather than
/// bubbling up to a whole-batch retry, which would re-COPY partitions that already
/// committed and duplicate rows.
async fn binary_parallel_pass(
&self,
row_slice: &[DataRow],
worker_count_option: Option<usize>,
) -> Result<u64> {
// 8 workers was simply the optimal number found in local testing. In the future this should
// be refactored to be an argument passed to the SDK, chose not to do that for the moment to avoid a
// breaking change
let worker_count = worker_count_option.unwrap_or(8);
let part = row_slice.len().div_ceil(worker_count).max(1);
// Tighter budget than the idempotent read default: a connection drop during a
// partition COMMIT is ambiguous, so each retry risks a double-write. Two retries
// (three attempts total) matches horizon-data-core's `insert_batch` policy.
let policy = RetryPolicy::default().with_max_retries(2);
let count_vec = try_join_all(row_slice.chunks(part).map(|chunk| {
retry(policy, move || {
Self::copy_binary_partition(&self.pool, chunk)
})
}))
.await?;
Ok(count_vec.into_iter().sum())
}
/// Run one binary `COPY ... FROM STDIN WITH (FORMAT binary)` for `rows` inside its
/// own transaction on a single connection from `pool`, returning the number of rows
/// `COPYed`. Builds the binary stream ourselves (file header, one tuple per row with
/// field count plus length-prefixed binary fields, and the -1 trailer), streaming it
/// to the `PgCopyIn` sink, then commits.
///
/// Wrapping the `COPY` in a transaction makes the per-partition [`retry`] in
/// `binary_parallel_pass` idempotent for server-side aborts: a failed attempt is
/// dropped without `commit`, rolling back any partial work so the retry re-COPYs from
/// a clean slate. The residual case is a connection drop during `COMMIT`, which is
/// ambiguous (the commit may have landed); a retry can then re-COPY the partition, so
/// `data_row` writes are at-least-once.
#[allow(
clippy::single_call_fn,
reason = "Used to breakup binary_parallel_pass function for readability"
)]
async fn copy_binary_partition(pool: &PgPool, row_slice: &[DataRow]) -> Result<u64> {
let mut transaction = pool.begin().await?;
let copied = {
let mut copy = transaction
.copy_in_raw(
"COPY horizon_public.data_row \
(data_stream_id, datetime, vector, data_type, specification_id, \
vector_start_bound, vector_end_bound, created_datetime, modified_datetime) \
FROM STDIN WITH (FORMAT binary)",
)
.await?;
// File header: 11-byte signature, int32 flags = 0, int32 header-extension length = 0.
let mut buf: Vec<u8> = Vec::with_capacity(16 * 1024 * 1024);
buf.extend_from_slice(b"PGCOPY\n\xFF\r\n\0");
buf.extend_from_slice(&0_i32.to_be_bytes()); // flags
buf.extend_from_slice(&0_i32.to_be_bytes()); // header extension area length
for row in row_slice {
let now = Utc::now();
buf.extend_from_slice(&9_i16.to_be_bytes()); // field count for this tuple
put_field(&mut buf, Some(row.data_stream_id.as_bytes()));
put_timestamptz(&mut buf, row.datetime);
put_float8_array_field(&mut buf, &row.vector);
put_field(&mut buf, Some(row.data_type.as_bytes()));
put_field(&mut buf, Some(row.specification_id.as_bytes()));
put_field(&mut buf, Some(&row.vector_start_bound.to_be_bytes()));
put_field(&mut buf, Some(&row.vector_end_bound.to_be_bytes()));
put_timestamptz(&mut buf, row.created_datetime.unwrap_or(now));
put_timestamptz(&mut buf, now);
// Stream in ~16 MB chunks so the whole payload never sits in memory at once.
if buf.len() >= 16 * 1024 * 1024 {
copy.send(buf.as_slice()).await?;
buf.clear();
}
}
// File trailer: int16 = -1 marks end of data.
buf.extend_from_slice(&(-1_i16).to_be_bytes());
copy.send(buf.as_slice()).await?;
copy.finish().await?
};
transaction.commit().await?;
Ok(copied)
}
/// If `created_datetime` is `None`, defaults to now (allows backfilling historical data).
pub async fn insert_data_row(&self, data_row: &DataRow) -> Result<DataRow> {
let now = Utc::now();
Ok(sqlx::query_as!(
DataRow,
r#"
INSERT INTO horizon_public.data_row
(data_stream_id, datetime, vector, data_type, specification_id,
vector_start_bound, vector_end_bound, created_datetime, modified_datetime)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
RETURNING created_datetime, modified_datetime, data_stream_id,
datetime, vector, data_type, specification_id,
vector_start_bound, vector_end_bound
"#,
data_row.data_stream_id,
data_row.datetime,
&data_row.vector,
data_row.data_type,
data_row.specification_id,
data_row.vector_start_bound,
data_row.vector_end_bound,
data_row.created_datetime.unwrap_or(now),
now,
)
.fetch_one(&self.pool)
.await?)
}
pub async fn insert_data_row_batch(&self, row_slice: &[DataRow]) -> Result<u64> {
self.binary_parallel_pass(row_slice, None).await
}
/// Ordered by datetime. Joins through `data_stream` for RLS scoping. Retried under
/// the default [`RetryPolicy`] on transient Postgres errors; reads are naturally
/// idempotent.
pub async fn list_data_rows(&self, data_stream_id: Uuid) -> Result<Vec<DataRow>> {
retry(RetryPolicy::default(), || async move {
Ok(sqlx::query_as!(
DataRow,
r#"
SELECT dr.created_datetime, dr.modified_datetime, dr.data_stream_id,
dr.datetime, dr.vector, dr.data_type, dr.specification_id,
dr.vector_start_bound, dr.vector_end_bound
FROM horizon_public.data_row dr
JOIN horizon_public.data_stream ds ON dr.data_stream_id = ds.id
WHERE dr.data_stream_id = $1
ORDER BY dr.datetime
"#,
data_stream_id,
)
.fetch_all(&self.pool)
.await?)
})
.await
}
}
/// Append one COPY-binary field: a 4-byte big-endian length prefix followed by
/// the field's binary bytes. `None` encodes a SQL NULL as a length of -1 with no
/// bytes following.
#[allow(
clippy::cast_possible_truncation,
clippy::cast_possible_wrap,
clippy::as_conversions,
reason = "Field lengths are bounded by PostgreSQL protocol limits, well within i32 range"
)]
fn put_field(buf: &mut Vec<u8>, bytes: Option<&[u8]>) {
match bytes {
Some(b) => {
buf.extend_from_slice(&(b.len() as i32).to_be_bytes());
buf.extend_from_slice(b);
}
None => buf.extend_from_slice(&(-1_i32).to_be_bytes()),
}
}
/// Encode `values` as a 1-D `float8[]` (no NULLs) directly into `buf`, including
/// the COPY field length prefix. Avoids a temporary allocation compared to building
/// a `Vec<u8>` and then copying it in via `put_field`.
#[allow(
clippy::single_call_fn,
reason = "Re-usable function that could be used for other array fields in the future"
)]
#[allow(
clippy::cast_possible_truncation,
clippy::cast_possible_wrap,
clippy::as_conversions,
clippy::arithmetic_side_effects,
reason = "Vector lengths are bounded by application constraints, well within i32 range"
)]
fn put_float8_array_field(buf: &mut Vec<u8>, values: &[f64]) {
const FLOAT8_OID: i32 = 701;
// 20-byte array header + 12 bytes per element (4-byte length + 8-byte f64).
let array_len = (20 + values.len() * 12) as i32;
buf.extend_from_slice(&array_len.to_be_bytes()); // COPY field length
buf.extend_from_slice(&1_i32.to_be_bytes()); // number of dimensions
buf.extend_from_slice(&0_i32.to_be_bytes()); // flags (no NULLs)
buf.extend_from_slice(&FLOAT8_OID.to_be_bytes()); // element OID
buf.extend_from_slice(&(values.len() as i32).to_be_bytes()); // dim length
buf.extend_from_slice(&1_i32.to_be_bytes()); // lower bound (1-based)
for &v in values {
buf.extend_from_slice(&8_i32.to_be_bytes()); // element byte length
buf.extend_from_slice(&v.to_be_bytes());
}
}
/// Append a non-null `timestamptz` field: micros since the PG epoch (2000-01-01)
/// as a big-endian `i64`, length-prefixed like any other field.
#[allow(
clippy::arithmetic_side_effects,
reason = "Timestamps are always after the PG epoch (2000-01-01); subtraction cannot underflow"
)]
fn put_timestamptz(buf: &mut Vec<u8>, dt: DateTime<Utc>) {
/// Microseconds between the Unix epoch (1970-01-01) and the `PostgreSQL` epoch
/// (2000-01-01) — timestamptz binary is micros since the PG epoch.
const PG_EPOCH_OFFSET_MICROS: i64 = 946_684_800_000_000;
let micros = dt.timestamp_micros() - PG_EPOCH_OFFSET_MICROS;
put_field(buf, Some(µs.to_be_bytes()));
}