use async_trait::async_trait;
use rust_decimal::Decimal;
use rust_decimal::prelude::ToPrimitive;
use sqlx::{PgPool, Row};
use time::{Duration, OffsetDateTime};
use tracing::{debug, info};
use crate::arrow::array::{
Array, ArrayRef, Decimal128Array, RecordBatch, StringArray, TimestampMicrosecondArray,
};
use crate::encode::schema::{
self, VALUE_PRECISION, VALUE_SCALE, VERSION_PRECISION, VERSION_SCALE, col,
};
use crate::error::{Error, Result};
use crate::planner::TimeRange;
use crate::tiering::store::{BatchStream, HotStore, PartitionId, ScanSpec};
use crate::watermark::TieringWatermark;
#[derive(Debug, Clone)]
pub struct PostgresHot {
pool: PgPool,
scan_chunk_rows: usize,
integrity_constraints: bool,
ddl_lock_timeout: Duration,
}
impl PostgresHot {
const DEFAULT_SCAN_CHUNK_ROWS: usize = crate::config::defaults::SCAN_CHUNK_ROWS;
const DEFAULT_DDL_LOCK_TIMEOUT: Duration = Duration::seconds(3);
pub fn new(pool: PgPool) -> Self {
Self {
pool,
scan_chunk_rows: Self::DEFAULT_SCAN_CHUNK_ROWS,
integrity_constraints: true,
ddl_lock_timeout: Self::DEFAULT_DDL_LOCK_TIMEOUT,
}
}
pub fn ddl_lock_timeout(mut self, timeout: Duration) -> Self {
self.ddl_lock_timeout = timeout.max(Duration::ZERO);
self
}
pub fn integrity_constraints(mut self, enabled: bool) -> Self {
self.integrity_constraints = enabled;
self
}
pub fn scan_chunk_rows(mut self, rows: usize) -> Self {
self.scan_chunk_rows = rows.max(1);
self
}
pub fn pool(&self) -> &PgPool {
&self.pool
}
pub async fn create_table(&self, table: &str) -> Result<()> {
self.create_table_with_key(
table,
&crate::encode::schema::MERGE_KEY
.iter()
.map(|s| (*s).to_string())
.collect::<Vec<_>>(),
&[],
crate::config::TimeModel::Interval,
)
.await
}
pub async fn create_table_with_key(
&self,
table: &str,
merge_key: &[String],
extra: &[crate::arrow::datatypes::Field],
time_model: crate::config::TimeModel,
) -> Result<()> {
let mut extra_ddl = String::new();
for f in extra {
let check = f
.metadata()
.get(crate::config::CHECK_VALUES_KEY)
.map(|vals| {
let codes: Vec<&str> = vals.split(',').collect();
format!(
" CONSTRAINT {constraint:?} CHECK ({col:?} IN ({list}))",
constraint = format!("{}_known", f.name()),
col = f.name(),
list = sql_code_list(&codes),
)
})
.or_else(|| {
crate::config::declared_value_check(f).map(|kind| {
format!(
" CONSTRAINT {constraint:?} CHECK ({col:?} ~ '{pattern}')",
constraint = format!("{}_shape", f.name()),
col = f.name(),
pattern = value_check_pattern(kind),
)
})
})
.unwrap_or_default();
extra_ddl.push_str(&format!(
"{:?} {} {}{},\n ",
f.name(),
pg_type(f.data_type())?,
if f.is_nullable() { "" } else { "NOT NULL" },
check,
));
}
let melo_ddl = match merge_key.iter().any(|c| c == col::MELO_ID) {
true => " NOT NULL",
false => "",
};
let to_ddl = match time_model.has_interval_end() {
true => {
" NOT NULL\n CONSTRAINT interval_forward CHECK (\"to\" > \"from\")"
}
false => "",
};
let ddl = format!(
r#"
CREATE TABLE IF NOT EXISTS "{table}" (
malo_id TEXT NOT NULL,
melo_id TEXT{melo_ddl},
-- Canonical OBIS only. This column is part of the merge key,
-- so two spellings of one channel would let a correction fail
-- to supersede the value it corrects. The canonical form omits
-- the storage group when it is unused (255), and keeps it when
-- it carries information, so `*255` is the one suffix that must
-- never appear. Failing the write beats resolving wrongly later.
obis_code TEXT NOT NULL
CONSTRAINT obis_code_canonical CHECK (
obis_code ~ '^[0-9]+-[0-9]+:[0-9]+\.[0-9]+\.[0-9]+(\*[0-9]+)?$'
AND obis_code !~ '\*255$'
),
-- Commodity and unit are checked against `metering`'s own code
-- lists, rendered below rather than spelled out here: a second
-- copy of the domain's vocabulary in DDL is a copy that drifts.
sparte TEXT NOT NULL
CONSTRAINT sparte_known CHECK (sparte IN ({sparte_codes})),
"from" TIMESTAMPTZ NOT NULL,
-- On an interval table: the span's exclusive end, present and
-- after the start. On a point table there is no end — a
-- Zählerstand is a register value at an instant — so the column
-- is null and `to IS NULL` is what tells a reader that `value`
-- is a cumulative reading rather than energy over a span.
"to" TIMESTAMPTZ{to_ddl},
value NUMERIC({VALUE_PRECISION},{VALUE_SCALE}) NOT NULL,
-- Water is m³ and gas may be either side of the Brennwert
-- conversion, so the number's dimension is stored, never implied
-- by the column name.
unit TEXT NOT NULL
CONSTRAINT unit_known CHECK (unit IN ({unit_codes})),
-- Quality is checked against `metering`'s own code list, rendered
-- below like sparte/unit: the stored value is the resolved reading's
-- quality, and a drifting literal must fail the write, not read back
-- as an unknown flag on the authoritative store.
quality TEXT NOT NULL
CONSTRAINT quality_known CHECK (quality IN ({quality_codes})),
resolution TEXT,
source_kind TEXT NOT NULL,
source_detail TEXT,
provenance TEXT,
version NUMERIC({VERSION_PRECISION},{VERSION_SCALE}) NOT NULL,
-- Canonical `<Marktpartner-ID>:<YYYY-MM>` only. The operator
-- half is a `metering::BdewCode` — thirteen digits, what MSCONS
-- carries in NAD+MS — so it cannot hold the separator that the
-- one-operator exclusion below reads it back with
-- (`split_part(version_scope, ':', 1)`).
--
-- The check digit is deliberately not checked: BDEW's
-- Anwendungshilfe §2.3 carves out GS1-issued GLNs, so a
-- well-formed Marktpartner-ID may legitimately fail it.
version_scope TEXT NOT NULL
CONSTRAINT version_scope_canonical CHECK (
version_scope ~ '^[0-9]{{13}}:[0-9]{{4}}-(0[1-9]|1[0-2])$'
),
recorded_at TIMESTAMPTZ NOT NULL,
-- The day this reading is balanced on: the Berlin calendar day,
-- or the 06:00-06:00 Gastag for gas. Derived by the encoder from
-- `from` and `sparte` and stored because no portable SQL
-- expresses the rule — see `encode::schema`. Not checked here:
-- PostgreSQL would need the Gastag rule to check it, which is
-- the very thing being avoided.
balancing_day DATE NOT NULL,
{extra_ddl}
PRIMARY KEY ({pk})
) PARTITION BY RANGE ("from")
"#,
extra_ddl = extra_ddl,
sparte_codes = sql_code_list(metering::Sparte::CODES),
unit_codes = sql_code_list(metering::interval::MeasurementUnit::CODES),
quality_codes = sql_code_list(metering::QualityFlag::CODES),
pk = merge_key
.iter()
.map(|c| format!("{c:?}"))
.chain(std::iter::once("version".to_string()))
.collect::<Vec<_>>()
.join(", "),
);
sqlx::query(&ddl).execute(&self.pool).await.map_err(pg)?;
self.verify_declaration(table, merge_key, time_model)
.await?;
let idx = format!(
r#"CREATE INDEX IF NOT EXISTS "{table}_malo_from_idx" ON "{table}" (malo_id, "from")"#
);
sqlx::query(&idx).execute(&self.pool).await.map_err(pg)?;
info!(table, "hot table ready");
Ok(())
}
async fn verify_declaration(
&self,
table: &str,
merge_key: &[String],
time_model: crate::config::TimeModel,
) -> Result<()> {
let mut conn = self.pool.acquire().await.map_err(pg)?;
let expected: Vec<String> = merge_key
.iter()
.cloned()
.chain(std::iter::once(col::VERSION.to_string()))
.collect();
let stored = primary_key_columns(&mut conn, table).await?;
if stored != expected {
return Err(Error::config(format!(
"{table} already exists with the primary key ({}), but this configuration \
declares the merge key ({}). The primary key is the merge key, and \
`CREATE TABLE IF NOT EXISTS` cannot change it — so resolution would \
partition by one key while the table enforced the other, and a reading \
the wider key exists to keep apart would be dropped with nothing \
reporting it. Restore the declaration, or create a new table",
stored.join(", "),
expected.join(", "),
)));
}
let stored_model = has_interval_end(&mut conn, table).await?;
if stored_model != time_model.has_interval_end() {
return Err(Error::config(format!(
"{table} already exists as a {} table, but this configuration declares {}. \
`value` is interval energy on one and a cumulative register reading on the \
other, so one table cannot hold both. Declare a second table",
match stored_model {
true => crate::config::TimeModel::Interval,
false => crate::config::TimeModel::Point,
},
time_model,
)));
}
Ok(())
}
}
async fn add_integrity_constraints(
conn: &mut sqlx::PgConnection,
table: &str,
partition: &str,
lock_timeout: Duration,
) -> Result<()> {
sqlx::query("CREATE EXTENSION IF NOT EXISTS btree_gist")
.execute(&mut *conn)
.await
.map_err(|e| {
Error::config(format!(
"overlap exclusion needs the btree_gist extension, and creating it \
failed ({e}). Install it as a superuser, or disable the check with \
PostgresHot::integrity_constraints(false) and accept that an \
overlapping delivery is then detected by completeness rather than \
refused"
))
})?;
let key = primary_key_columns(&mut *conn, table).await?;
let equality: Vec<String> = key
.iter()
.filter(|c| c.as_str() != col::FROM)
.map(|c| format!("{c:?} WITH ="))
.collect();
if equality.is_empty() {
return Err(Error::config(format!(
"{table} has no primary key columns besides {:?}, so an overlap \
exclusion would compare every row against every other",
col::FROM
)));
}
if has_interval_end(&mut *conn, table).await? {
let ddl = format!(
r#"ALTER TABLE "{partition}" ADD CONSTRAINT "{partition}_no_overlap"
EXCLUDE USING gist ({}, tstzrange("from", "to", '[)') WITH &&)"#,
equality.join(", "),
);
sqlx::query(&ddl)
.execute(&mut *conn)
.await
.map_err(|e| pg_ddl(partition, "add overlap exclusion", lock_timeout, e))?;
}
let key = primary_key_columns(&mut *conn, table).await?;
let merge_key: Vec<String> = key
.iter()
.filter(|c| c.as_str() != col::VERSION)
.map(|c| format!("{c:?} WITH ="))
.collect();
let ddl = format!(
r#"ALTER TABLE "{partition}" ADD CONSTRAINT "{partition}_one_operator"
EXCLUDE USING gist ({}, split_part(version_scope, ':', 1) WITH <>)"#,
merge_key.join(", "),
);
sqlx::query(&ddl)
.execute(&mut *conn)
.await
.map_err(|e| pg_ddl(partition, "add operator exclusion", lock_timeout, e))?;
Ok(())
}
async fn has_interval_end(conn: &mut sqlx::PgConnection, table: &str) -> Result<bool> {
sqlx::query_scalar::<_, bool>(
r#"SELECT a.attnotnull
FROM pg_attribute a
WHERE a.attrelid = $1::regclass
AND a.attname = 'to'
AND NOT a.attisdropped"#,
)
.bind(format!("\"{table}\""))
.fetch_optional(&mut *conn)
.await
.map_err(pg)?
.ok_or_else(|| {
Error::config(format!(
"{table} has no {:?} column, so its time model cannot be read",
col::TO
))
})
}
async fn primary_key_columns(conn: &mut sqlx::PgConnection, table: &str) -> Result<Vec<String>> {
let rows: Vec<String> = sqlx::query_scalar(
r#"SELECT a.attname::text
FROM pg_index i
JOIN pg_attribute a
ON a.attrelid = i.indrelid
AND a.attnum = ANY(i.indkey)
WHERE i.indrelid = $1::regclass
AND i.indisprimary
ORDER BY array_position(i.indkey, a.attnum)"#,
)
.bind(format!("\"{table}\""))
.fetch_all(&mut *conn)
.await
.map_err(pg)?;
if rows.is_empty() {
return Err(Error::config(format!(
"{table} has no primary key; the overlap exclusion derives its \
equality columns from it"
)));
}
Ok(rows)
}
impl PostgresHot {
async fn insert_reporting(
&self,
table: &str,
merge_key: &[String],
batch: &RecordBatch,
) -> Result<Vec<crate::session::Displacement>> {
use crate::session::{Displacement, Effect, StoredValue};
let n = batch.num_rows();
if n == 0 {
return Ok(Vec::new());
}
let identity: Vec<String> = merge_key
.iter()
.filter(|c| !crate::encode::schema::MERGE_KEY.contains(&c.as_str()))
.cloned()
.collect();
let value_of = |row: usize| -> Result<StoredValue> {
let r = RowView::new(batch, row)?;
Ok(StoredValue {
value: r.value,
unit: r
.unit
.parse()
.map_err(|e| Error::decode(col::UNIT, format!("{:?}: {e}", r.unit)))?,
quality: r
.quality
.parse()
.map_err(|e| Error::decode(col::QUALITY, format!("{:?}: {e}", r.quality)))?,
version: crate::version::ScopedVersion::new(
crate::version::VersionScope::parse(r.version_scope)?,
decimal_to_version(r.version)?,
),
recorded_at: r.recorded_at,
})
};
let mut rows = Vec::with_capacity(n);
for row in 0..n {
let r = RowView::new(batch, row)?;
let mut ident = Vec::with_capacity(identity.len());
for name in &identity {
ident.push((name.clone(), key_column(batch, name, row)?.to_string()));
}
rows.push((
(r.malo.to_string(), r.obis.to_string(), r.from, ident),
value_of(row)?,
r.to,
));
}
let mut tx = self.pool.begin().await.map_err(pg)?;
let malo: Vec<String> = rows.iter().map(|(k, ..)| k.0.clone()).collect();
let obis: Vec<String> = rows.iter().map(|(k, ..)| k.1.clone()).collect();
let from: Vec<OffsetDateTime> = rows.iter().map(|(k, ..)| k.2).collect();
let ident_select = identity
.iter()
.map(|c| format!(", s.{c:?}"))
.collect::<String>();
let sql = format!(
r#"SELECT s.malo_id, s.obis_code, s."from", s.value, s.unit, s.quality,
s.version, s.version_scope, s.recorded_at{ident_select}
FROM "{table}" s
JOIN unnest($1::text[], $2::text[], $3::timestamptz[])
AS k(malo_id, obis_code, "from")
ON s.malo_id = k.malo_id
AND s.obis_code = k.obis_code
AND s."from" = k."from""#
);
let stored = sqlx::query(&sql)
.bind(&malo)
.bind(&obis)
.bind(&from)
.fetch_all(&mut *tx)
.await
.map_err(pg)?;
type Key = (String, String, OffsetDateTime, Vec<(String, String)>);
let mut current: std::collections::HashMap<Key, StoredValue> = Default::default();
let mut seen: std::collections::HashSet<(Key, u128)> = Default::default();
for row in &stored {
let mut ident = Vec::with_capacity(identity.len());
for (i, name) in identity.iter().enumerate() {
ident.push((name.clone(), row.try_get::<String, _>(9 + i).map_err(pg)?));
}
let key: Key = (
row.try_get(0).map_err(pg)?,
row.try_get(1).map_err(pg)?,
row.try_get(2).map_err(pg)?,
ident,
);
let unit: String = row.try_get(4).map_err(pg)?;
let quality: String = row.try_get(5).map_err(pg)?;
let scope: String = row.try_get(7).map_err(pg)?;
let held = StoredValue {
value: row.try_get(3).map_err(pg)?,
unit: unit
.parse()
.map_err(|e| Error::decode(col::UNIT, format!("{unit:?}: {e}")))?,
quality: quality
.parse()
.map_err(|e| Error::decode(col::QUALITY, format!("{quality:?}: {e}")))?,
version: crate::version::ScopedVersion::new(
crate::version::VersionScope::parse(&scope)?,
decimal_to_version(row.try_get(6).map_err(pg)?)?,
),
recorded_at: row.try_get(8).map_err(pg)?,
};
seen.insert((key.clone(), held.version.version().get()));
match current.get(&key) {
Some(best) if best.version.try_cmp(&held.version)? != std::cmp::Ordering::Less => {}
_ => {
current.insert(key, held);
}
}
}
self.insert_rows(&mut tx, table, merge_key, batch).await?;
tx.commit().await.map_err(pg)?;
let mut order: Vec<usize> = (0..rows.len()).collect();
order.sort_by_key(|&i| rows[i].1.version.version().get());
let mut out = vec![None; rows.len()];
for i in order {
let (key, written, to) = &rows[i];
let prior = current.get(key).cloned();
let replayed = seen.contains(&(key.clone(), written.version.version().get()));
let effect = match &prior {
None => Effect::Inserted,
Some(_) if replayed => Effect::Duplicate,
Some(p) => match p.version.try_cmp(&written.version)? {
std::cmp::Ordering::Less => Effect::Superseded,
_ => Effect::Shadowed,
},
};
if effect.changed_current_value() {
current.insert(key.clone(), written.clone());
}
out[i] = Some(Displacement {
malo_id: key.0.clone(),
obis_code: key.1.clone(),
from: key.2,
to: *to,
identity: key.3.clone(),
effect,
superseded: prior,
written: written.clone(),
});
}
Ok(out.into_iter().flatten().collect())
}
async fn append_batch(
&self,
table: &str,
merge_key: &[String],
batch: &RecordBatch,
) -> Result<u64> {
let mut conn = self.pool.acquire().await.map_err(pg)?;
self.insert_rows(&mut conn, table, merge_key, batch).await
}
async fn insert_rows(
&self,
conn: &mut sqlx::PgConnection,
table: &str,
merge_key: &[String],
batch: &RecordBatch,
) -> Result<u64> {
let core = crate::encode::schema::storage_schema(&[]);
let extra: Vec<String> = batch
.schema()
.fields()
.iter()
.filter(|f| core.field_with_name(f.name()).is_err())
.map(|f| f.name().clone())
.collect();
let conflict = merge_key
.iter()
.map(|c| format!("{c:?}"))
.chain(std::iter::once("version".to_string()))
.collect::<Vec<_>>()
.join(", ");
let n = batch.num_rows();
let mut malo = Vec::with_capacity(n);
let mut melo: Vec<Option<String>> = Vec::with_capacity(n);
let mut obis = Vec::with_capacity(n);
let mut sparte = Vec::with_capacity(n);
let mut from = Vec::with_capacity(n);
let mut to: Vec<Option<OffsetDateTime>> = Vec::with_capacity(n);
let mut value = Vec::with_capacity(n);
let mut unit = Vec::with_capacity(n);
let mut quality = Vec::with_capacity(n);
let mut resolution: Vec<Option<String>> = Vec::with_capacity(n);
let mut source_kind = Vec::with_capacity(n);
let mut source_detail: Vec<Option<String>> = Vec::with_capacity(n);
let mut provenance: Vec<Option<String>> = Vec::with_capacity(n);
let mut version = Vec::with_capacity(n);
let mut version_scope = Vec::with_capacity(n);
let mut recorded_at = Vec::with_capacity(n);
let mut balancing_day = Vec::with_capacity(n);
for row in 0..n {
let r = RowView::new(batch, row)?;
malo.push(r.malo.to_string());
melo.push(r.melo.map(str::to_string));
obis.push(r.obis.to_string());
sparte.push(r.sparte.to_string());
from.push(r.from);
to.push(r.to);
value.push(r.value);
unit.push(r.unit.to_string());
quality.push(r.quality.to_string());
resolution.push(r.resolution.map(str::to_string));
source_kind.push(r.source_kind.to_string());
source_detail.push(r.source_detail.map(str::to_string));
provenance.push(r.provenance.map(str::to_string));
version.push(r.version);
version_scope.push(r.version_scope.to_string());
recorded_at.push(r.recorded_at);
balancing_day.push(r.balancing_day);
}
let mut extra_values: Vec<Vec<Option<String>>> = Vec::with_capacity(extra.len());
for name in &extra {
let column = batch
.column_by_name(name)
.ok_or_else(|| Error::encode(name, "declared column missing from batch"))?;
let strings = column
.as_any()
.downcast_ref::<StringArray>()
.ok_or_else(|| Error::encode(name, "extra columns must be text"))?;
extra_values.push(
(0..strings.len())
.map(|i| (!strings.is_null(i)).then(|| strings.value(i).to_string()))
.collect(),
);
}
let extra_cols = extra.iter().map(|c| format!(", {c:?}")).collect::<String>();
let extra_params = (0..extra.len())
.map(|i| format!(", ${}::text[]", core_column_count() + 1 + i))
.collect::<String>();
let sql = format!(
r#"INSERT INTO "{table}" ({core_cols}{extra_cols})
SELECT * FROM unnest(
$1::text[], $2::text[], $3::text[], $4::text[],
$5::timestamptz[], $6::timestamptz[], $7::numeric[], $8::text[],
$9::text[], $10::text[], $11::text[], $12::text[], $13::text[],
$14::numeric[], $15::text[], $16::timestamptz[], $17::date[]{extra_params}
)
ON CONFLICT ({conflict}) DO NOTHING"#,
core_cols = scan_columns(),
);
let mut query = sqlx::query(&sql)
.bind(&malo)
.bind(&melo)
.bind(&obis)
.bind(&sparte)
.bind(&from)
.bind(&to)
.bind(&value)
.bind(&unit)
.bind(&quality)
.bind(&resolution)
.bind(&source_kind)
.bind(&source_detail)
.bind(&provenance)
.bind(&version)
.bind(&version_scope)
.bind(&recorded_at)
.bind(&balancing_day);
for values in &extra_values {
query = query.bind(values);
}
let inserted = query.execute(&mut *conn).await.map_err(pg)?.rows_affected();
let skipped = (n as u64).saturating_sub(inserted);
let metrics = crate::observe::metrics();
let attrs = crate::observe::table(table);
metrics.rows_written.add(inserted, &attrs);
metrics.rows_deduplicated.add(skipped, &attrs);
if skipped > 0 {
let identity: Vec<&String> = merge_key
.iter()
.filter(|c| {
!crate::encode::schema::MERGE_KEY.contains(&c.as_str())
&& c.as_str() != col::MELO_ID
})
.collect();
let join = merge_key
.iter()
.map(|c| format!(r#"stored.{c:?} = incoming.{c:?}"#))
.chain(std::iter::once(
"stored.version = incoming.version".to_string(),
))
.collect::<Vec<_>>()
.join(" AND ");
let identity_params = (0..identity.len())
.map(|i| format!(", ${}::text[]", 7 + i))
.collect::<String>();
let identity_cols = identity
.iter()
.map(|c| format!(", {c:?}"))
.collect::<String>();
let sql = format!(
r#"SELECT
count(*) FILTER (
WHERE stored.value IS DISTINCT FROM incoming.value) AS restated,
count(*) FILTER (
WHERE stored.melo_id IS DISTINCT FROM incoming.melo_id) AS relocated
FROM unnest(
$1::text[], $2::text[], $3::timestamptz[],
$4::numeric[], $5::numeric[], $6::text[]{identity_params}
) AS incoming(malo_id, obis_code, "from", version, value, melo_id{identity_cols})
JOIN "{table}" stored ON {join}"#
);
let mut query = sqlx::query_as::<_, (i64, i64)>(&sql)
.bind(&malo)
.bind(&obis)
.bind(&from)
.bind(&version)
.bind(&value)
.bind(&melo);
for name in &identity {
let index = extra
.iter()
.position(|e| e == *name)
.ok_or_else(|| Error::encode(name.as_str(), "identity column missing"))?;
query = query.bind(&extra_values[index]);
}
let (restated, relocated) = query.fetch_one(&mut *conn).await.map_err(pg)?;
if restated > 0 {
return Err(Error::IntegrityViolation {
table: table.to_string(),
constraint: Some("version_identifies_one_assertion".to_string()),
detail: format!(
"{restated} row(s) restate a different value under an existing \
version — a version identifies one assertion, so a corrected \
value needs a higher version"
),
});
}
if relocated > 0 {
return Err(Error::IntegrityViolation {
table: table.to_string(),
constraint: Some("melo_identifies_the_reading".to_string()),
detail: format!(
"{relocated} row(s) name a different Messlokation than the row \
already stored for the same reading. A Marktlokation may be \
measured by several Messlokationen — a Mehrfamilienhaus, a house \
with an Einliegerwohnung — and this table does not identify a \
reading by its, so two meters' registers share a merge key and \
one of them is dropped. Declare \
TableConfig::identify_by_melo(true)"
),
});
}
debug!(table, skipped, "replayed rows already present");
}
Ok(inserted)
}
fn chunked_scan(&self, relation: &str, range: TimeRange, spec: &ScanSpec) -> BatchStream {
let pool = self.pool.clone();
let relation = relation.to_string();
let chunk = spec.chunk_rows().unwrap_or(self.scan_chunk_rows);
let extra: Vec<String> = spec.extra().to_vec();
let extra_select = extra.iter().map(|c| format!(", {c:?}")).collect::<String>();
let cursor_columns = spec.cursor_columns();
let order_by = cursor_columns
.iter()
.map(|c| format!("{c:?}"))
.collect::<Vec<_>>()
.join(", ");
let cursor_positions: Vec<usize> = cursor_columns
.iter()
.map(|c| projection_index(c, &extra))
.collect::<Result<Vec<_>>>()
.unwrap_or_default();
Box::pin(async_stream::try_stream! {
if cursor_positions.len() != cursor_columns.len() {
Err(Error::config(format!(
"cursor columns {cursor_columns:?} are not all projected by the scan"
)))?;
}
let mut cursor: Option<Vec<CursorValue>> = None;
loop {
let mut sql =
format!(
r#"SELECT {}{extra_select} FROM "{relation}" WHERE true"#,
scan_columns()
);
let mut n = 0;
if range.start().is_some() {
n += 1;
sql.push_str(&format!(r#" AND "from" >= ${n}"#));
}
if range.end().is_some() {
n += 1;
sql.push_str(&format!(r#" AND "from" < ${n}"#));
}
if cursor.is_some() {
let placeholders = (0..cursor_columns.len())
.map(|i| format!("${}", n + 1 + i))
.collect::<Vec<_>>()
.join(", ");
sql.push_str(&format!(r#" AND ({order_by}) > ({placeholders})"#));
}
sql.push_str(&format!(r#" ORDER BY {order_by} LIMIT {chunk}"#));
let mut query = sqlx::query(&sql);
if let Some(start) = range.start() {
query = query.bind(start);
}
if let Some(end) = range.end() {
query = query.bind(end);
}
if let Some(values) = &cursor {
for value in values {
query = match value {
CursorValue::Text(v) => query.bind(v.clone()),
CursorValue::Timestamp(v) => query.bind(*v),
CursorValue::Numeric(v) => query.bind(*v),
};
}
}
let rows = query.fetch_all(&pool).await.map_err(pg)?;
if rows.is_empty() {
break;
}
let last = rows.last().expect("non-empty");
cursor = Some(
cursor_positions
.iter()
.map(|i| CursorValue::read(last, *i))
.collect::<Result<Vec<_>>>()?,
);
let exhausted = rows.len() < chunk;
for batch in rows_to_batches(rows, &extra)? {
yield batch;
}
if exhausted {
break;
}
}
})
}
async fn create_partition(
&self,
table: &str,
id: &PartitionId,
step: Duration,
) -> Result<bool> {
let name = id.relation_name()?;
if self.relation_exists(&name).await? {
return Ok(false);
}
let mut tx = self.begin_ddl().await?;
sqlx::query("SELECT pg_advisory_xact_lock($1)")
.bind(lock_key("partition", &name))
.execute(&mut *tx)
.await
.map_err(pg)?;
if relation_exists_in(&mut tx, &name).await? {
return Ok(false);
}
let end = id.start() + step;
let (lower, upper) = (pg_timestamp(id.start())?, pg_timestamp(end)?);
let timeout = self.ddl_lock_timeout;
let ddl = |op: &'static str, sql: String| (op, sql);
for (operation, sql) in [
ddl(
"create partition",
format!(
r#"CREATE TABLE "{name}" (
LIKE "{table}" INCLUDING DEFAULTS INCLUDING CONSTRAINTS
INCLUDING STORAGE INCLUDING COMMENTS
)"#
),
),
ddl(
"constrain partition",
format!(
r#"ALTER TABLE "{name}" ADD CONSTRAINT "{name}_bound"
CHECK ("from" >= '{lower}' AND "from" < '{upper}')"#
),
),
] {
sqlx::query(&sql)
.execute(&mut *tx)
.await
.map_err(|e| pg_ddl(&name, operation, timeout, e))?;
}
if self.integrity_constraints {
add_integrity_constraints(&mut tx, table, &name, timeout).await?;
}
sqlx::query(&format!(
r#"ALTER TABLE "{table}" ATTACH PARTITION "{name}"
FOR VALUES FROM ('{lower}') TO ('{upper}')"#
))
.execute(&mut *tx)
.await
.map_err(|e| pg_ddl(table, "attach partition", timeout, e))?;
sqlx::query(&format!(
r#"ALTER TABLE "{name}" DROP CONSTRAINT "{name}_bound""#
))
.execute(&mut *tx)
.await
.map_err(|e| pg_ddl(&name, "drop bound constraint", timeout, e))?;
tx.commit().await.map_err(pg)?;
debug!(table, partition = %name, "created hot partition");
Ok(true)
}
async fn begin_ddl(&self) -> Result<sqlx::Transaction<'_, sqlx::Postgres>> {
let mut tx = self.pool.begin().await.map_err(pg)?;
if self.ddl_lock_timeout > Duration::ZERO {
let ms = self.ddl_lock_timeout.whole_milliseconds().max(1);
sqlx::query(&format!("SET LOCAL lock_timeout = {ms}"))
.execute(&mut *tx)
.await
.map_err(pg)?;
}
Ok(tx)
}
async fn relation_exists(&self, name: &str) -> Result<bool> {
let row = sqlx::query_scalar::<_, bool>(
"SELECT EXISTS (SELECT 1 FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relname = $1 AND n.nspname = current_schema())",
)
.bind(name)
.fetch_one(&self.pool)
.await
.map_err(pg)?;
Ok(row)
}
}
const COLD_APPEND_LEASE_ATTEMPTS: u32 = 22;
const COLD_APPEND_LEASE_BUDGET_MS: u64 = {
let mut total = 0;
let mut attempt = 0;
while attempt < COLD_APPEND_LEASE_ATTEMPTS {
let step = if attempt < 5 { attempt } else { 5 };
total += 25u64 << step;
attempt += 1;
}
total
};
fn lock_key(purpose: &str, name: &str) -> i64 {
let mut hash: u64 = 0xcbf2_9ce4_8422_2325;
for byte in b"meterstore."
.iter()
.chain(purpose.as_bytes())
.chain(b":")
.chain(name.as_bytes())
{
hash ^= u64::from(*byte);
hash = hash.wrapping_mul(0x0000_0100_0000_01b3);
}
hash as i64
}
async fn relation_exists_in(conn: &mut sqlx::PgConnection, name: &str) -> Result<bool> {
sqlx::query_scalar::<_, bool>(
"SELECT EXISTS (SELECT 1 FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relname = $1 AND n.nspname = current_schema())",
)
.bind(name)
.fetch_one(&mut *conn)
.await
.map_err(pg)
}
struct PgTableLease {
connection: sqlx::pool::PoolConnection<sqlx::Postgres>,
key: i64,
table: String,
purpose: &'static str,
}
impl std::fmt::Debug for PgTableLease {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PgTableLease")
.field("table", &self.table)
.field("purpose", &self.purpose)
.finish_non_exhaustive()
}
}
#[async_trait]
impl crate::tiering::store::TableLease for PgTableLease {
async fn release(mut self: Box<Self>) -> Result<()> {
sqlx::query("SELECT pg_advisory_unlock($1)")
.bind(self.key)
.execute(&mut *self.connection)
.await
.map_err(pg)?;
debug!(table = %self.table, purpose = self.purpose, "lease released");
Ok(())
}
}
fn decimal_to_version(value: Decimal) -> Result<crate::version::Version> {
let digits = value.trunc().to_string();
crate::version::Version::new(
digits
.parse::<u128>()
.map_err(|e| Error::decode(col::VERSION, format!("{digits}: {e}")))?,
)
}
fn key_column<'a>(batch: &'a RecordBatch, name: &str, row: usize) -> Result<&'a str> {
let column = batch
.column_by_name(name)
.and_then(|c| c.as_any().downcast_ref::<StringArray>())
.ok_or_else(|| Error::encode(name, "expected a string column"))?;
if column.is_null(row) {
return Err(Error::encode(
name,
"is part of this table's merge key and must not be null: a null cannot \
identify a reading, and in SQL it does not compare equal to itself",
));
}
Ok(column.value(row))
}
fn value_check_pattern(kind: &str) -> &'static str {
crate::config::ValueCheck::from_code(kind)
.map_or("$^", crate::config::ValueCheck::shape_pattern)
}
fn sql_code_list(codes: &[&str]) -> String {
codes
.iter()
.map(|c| format!("'{}'", c.replace('\'', "''")))
.collect::<Vec<_>>()
.join(", ")
}
fn pg_type(ty: &crate::arrow::datatypes::DataType) -> Result<&'static str> {
use crate::arrow::datatypes::DataType;
Ok(match ty {
DataType::Utf8 | DataType::LargeUtf8 => "TEXT",
DataType::Boolean => "BOOLEAN",
DataType::Int16 => "SMALLINT",
DataType::Int32 => "INTEGER",
DataType::Int64 => "BIGINT",
DataType::Float64 => "DOUBLE PRECISION",
DataType::Date32 => "DATE",
DataType::Timestamp(_, _) => "TIMESTAMPTZ",
other => {
return Err(Error::config(format!(
"no PostgreSQL mapping for extra column type {other:?}"
)));
}
})
}
#[derive(Debug, Clone)]
enum CursorValue {
Text(String),
Timestamp(OffsetDateTime),
Numeric(Decimal),
}
impl CursorValue {
fn read(row: &sqlx::postgres::PgRow, index: usize) -> Result<Self> {
use crate::arrow::datatypes::DataType;
let core = schema::storage_schema(&[]);
let kind = core.fields().get(index).map(|f| f.data_type().clone());
Ok(match kind {
Some(DataType::Timestamp(_, _)) => {
Self::Timestamp(row.try_get::<OffsetDateTime, _>(index).map_err(pg)?)
}
Some(DataType::Decimal128(_, _)) => {
Self::Numeric(row.try_get::<Decimal, _>(index).map_err(pg)?)
}
_ => Self::Text(row.try_get::<String, _>(index).map_err(pg)?),
})
}
}
fn projection_index(column: &str, extra: &[String]) -> Result<usize> {
if let Ok(index) = crate::encode::schema::storage_schema(&[]).index_of(column) {
return Ok(index);
}
extra
.iter()
.position(|c| c == column)
.map(|i| core_column_count() + i)
.ok_or_else(|| Error::config(format!("column {column:?} is not projected by the scan")))
}
fn core_column_count() -> usize {
schema::storage_schema(&[]).fields().len()
}
fn scan_columns() -> String {
schema::storage_schema(&[])
.fields()
.iter()
.map(|f| format!("{:?}", f.name()))
.collect::<Vec<_>>()
.join(", ")
}
fn rows_to_batches(rows: Vec<sqlx::postgres::PgRow>, extra: &[String]) -> Result<Vec<RecordBatch>> {
if rows.is_empty() {
return Ok(Vec::new());
}
let n = rows.len();
let mut malo = Vec::with_capacity(n);
let mut melo: Vec<Option<String>> = Vec::with_capacity(n);
let mut obis = Vec::with_capacity(n);
let mut sparte = Vec::with_capacity(n);
let mut from = Vec::with_capacity(n);
let mut to: Vec<Option<i64>> = Vec::with_capacity(n);
let mut value = Vec::with_capacity(n);
let mut unit = Vec::with_capacity(n);
let mut quality = Vec::with_capacity(n);
let mut resolution: Vec<Option<String>> = Vec::with_capacity(n);
let mut source_kind = Vec::with_capacity(n);
let mut source_detail: Vec<Option<String>> = Vec::with_capacity(n);
let mut provenance: Vec<Option<String>> = Vec::with_capacity(n);
let mut version = Vec::with_capacity(n);
let mut version_scope = Vec::with_capacity(n);
let mut recorded_at = Vec::with_capacity(n);
let mut balancing_day = Vec::with_capacity(n);
for row in &rows {
malo.push(row.try_get::<String, _>(0).map_err(pg)?);
melo.push(row.try_get::<Option<String>, _>(1).map_err(pg)?);
obis.push(row.try_get::<String, _>(2).map_err(pg)?);
sparte.push(row.try_get::<String, _>(3).map_err(pg)?);
from.push(schema::micros(
row.try_get::<OffsetDateTime, _>(4).map_err(pg)?,
));
to.push(
row.try_get::<Option<OffsetDateTime>, _>(5)
.map_err(pg)?
.map(schema::micros),
);
value.push(decimal_to_i128(
row.try_get::<Decimal, _>(6).map_err(pg)?,
VALUE_SCALE,
col::VALUE,
)?);
unit.push(row.try_get::<String, _>(7).map_err(pg)?);
quality.push(row.try_get::<String, _>(8).map_err(pg)?);
resolution.push(row.try_get::<Option<String>, _>(9).map_err(pg)?);
source_kind.push(row.try_get::<String, _>(10).map_err(pg)?);
source_detail.push(row.try_get::<Option<String>, _>(11).map_err(pg)?);
provenance.push(row.try_get::<Option<String>, _>(12).map_err(pg)?);
version.push(decimal_to_i128(
row.try_get::<Decimal, _>(13).map_err(pg)?,
VERSION_SCALE,
col::VERSION,
)?);
version_scope.push(row.try_get::<String, _>(14).map_err(pg)?);
recorded_at.push(schema::micros(
row.try_get::<OffsetDateTime, _>(15).map_err(pg)?,
));
balancing_day.push(schema::date32(
row.try_get::<time::Date, _>(16).map_err(pg)?,
));
}
let tz: std::sync::Arc<str> = "UTC".into();
let columns: Vec<ArrayRef> = vec![
std::sync::Arc::new(StringArray::from(malo)),
std::sync::Arc::new(StringArray::from(melo)),
std::sync::Arc::new(StringArray::from(obis)),
std::sync::Arc::new(StringArray::from(sparte)),
std::sync::Arc::new(TimestampMicrosecondArray::from(from).with_timezone(tz.clone())),
std::sync::Arc::new(TimestampMicrosecondArray::from(to).with_timezone(tz.clone())),
std::sync::Arc::new(
Decimal128Array::from(value).with_precision_and_scale(VALUE_PRECISION, VALUE_SCALE)?,
),
std::sync::Arc::new(StringArray::from(unit)),
std::sync::Arc::new(StringArray::from(quality)),
std::sync::Arc::new(StringArray::from(resolution)),
std::sync::Arc::new(StringArray::from(source_kind)),
std::sync::Arc::new(StringArray::from(source_detail)),
std::sync::Arc::new(StringArray::from(provenance)),
std::sync::Arc::new(
Decimal128Array::from(version)
.with_precision_and_scale(VERSION_PRECISION, VERSION_SCALE)?,
),
std::sync::Arc::new(StringArray::from(version_scope)),
std::sync::Arc::new(TimestampMicrosecondArray::from(recorded_at).with_timezone(tz)),
std::sync::Arc::new(crate::arrow::array::Date32Array::from(balancing_day)),
];
let mut columns = columns;
let mut fields = Vec::with_capacity(extra.len());
for (i, name) in extra.iter().enumerate() {
let values: Vec<Option<String>> = rows
.iter()
.map(|r| r.try_get::<Option<String>, _>(core_column_count() + i))
.collect::<std::result::Result<_, _>>()
.map_err(pg)?;
columns.push(std::sync::Arc::new(StringArray::from(values)));
fields.push(crate::arrow::datatypes::Field::new(
name,
crate::arrow::datatypes::DataType::Utf8,
true,
));
}
Ok(vec![RecordBatch::try_new(
schema::storage_schema(&fields),
columns,
)?])
}
struct RowView<'a> {
malo: &'a str,
melo: Option<&'a str>,
obis: &'a str,
sparte: &'a str,
from: OffsetDateTime,
to: Option<OffsetDateTime>,
value: Decimal,
unit: &'a str,
quality: &'a str,
resolution: Option<&'a str>,
source_kind: &'a str,
source_detail: Option<&'a str>,
provenance: Option<&'a str>,
version: Decimal,
version_scope: &'a str,
recorded_at: OffsetDateTime,
balancing_day: time::Date,
}
impl<'a> RowView<'a> {
fn new(batch: &'a RecordBatch, row: usize) -> Result<Self> {
use crate::arrow::array::{Array, Decimal128Array, StringArray, TimestampMicrosecondArray};
fn text<'b>(batch: &'b RecordBatch, name: &str, row: usize) -> Result<&'b str> {
Ok(batch
.column_by_name(name)
.and_then(|c| c.as_any().downcast_ref::<StringArray>())
.ok_or_else(|| Error::encode(name, "expected a string column"))?
.value(row))
}
fn text_opt<'b>(batch: &'b RecordBatch, name: &str, row: usize) -> Option<&'b str> {
batch
.column_by_name(name)
.and_then(|c| c.as_any().downcast_ref::<StringArray>())
.filter(|a| !a.is_null(row))
.map(|a| a.value(row))
}
fn ts_opt(batch: &RecordBatch, name: &str, row: usize) -> Result<Option<OffsetDateTime>> {
let column = batch
.column_by_name(name)
.and_then(|c| c.as_any().downcast_ref::<TimestampMicrosecondArray>())
.ok_or_else(|| Error::encode(name, "expected a timestamp column"))?;
if column.is_null(row) {
return Ok(None);
}
OffsetDateTime::from_unix_timestamp_nanos(i128::from(column.value(row)) * 1_000)
.map(Some)
.map_err(|e| Error::encode(name, e.to_string()))
}
fn ts(batch: &RecordBatch, name: &str, row: usize) -> Result<OffsetDateTime> {
let micros = batch
.column_by_name(name)
.and_then(|c| c.as_any().downcast_ref::<TimestampMicrosecondArray>())
.ok_or_else(|| Error::encode(name, "expected a timestamp column"))?
.value(row);
OffsetDateTime::from_unix_timestamp_nanos(i128::from(micros) * 1_000)
.map_err(|e| Error::encode(name, e.to_string()))
}
fn date(batch: &RecordBatch, name: &str, row: usize) -> Result<time::Date> {
let days = batch
.column_by_name(name)
.and_then(|c| {
c.as_any()
.downcast_ref::<crate::arrow::array::Date32Array>()
})
.ok_or_else(|| Error::encode(name, "expected a date column"))?
.value(row);
schema::date_of(days)
}
fn dec(batch: &RecordBatch, name: &str, row: usize, scale: i8) -> Result<Decimal> {
let raw = batch
.column_by_name(name)
.and_then(|c| c.as_any().downcast_ref::<Decimal128Array>())
.ok_or_else(|| Error::encode(name, "expected a decimal column"))?
.value(row);
Decimal::try_from_i128_with_scale(raw, u32::try_from(scale).unwrap_or(0))
.map_err(|e| Error::encode(name, format!("{raw}: {e}")))
}
Ok(Self {
malo: text(batch, col::MALO_ID, row)?,
melo: text_opt(batch, col::MELO_ID, row),
obis: text(batch, col::OBIS_CODE, row)?,
sparte: text(batch, col::SPARTE, row)?,
from: ts(batch, col::FROM, row)?,
to: ts_opt(batch, col::TO, row)?,
value: dec(batch, col::VALUE, row, VALUE_SCALE)?,
unit: text(batch, col::UNIT, row)?,
quality: text(batch, col::QUALITY, row)?,
resolution: text_opt(batch, col::RESOLUTION, row),
source_kind: text(batch, col::SOURCE_KIND, row)?,
source_detail: text_opt(batch, col::SOURCE_DETAIL, row),
provenance: text_opt(batch, col::PROVENANCE, row),
version: dec(batch, col::VERSION, row, VERSION_SCALE)?,
version_scope: text(batch, col::VERSION_SCOPE, row)?,
recorded_at: ts(batch, col::RECORDED_AT, row)?,
balancing_day: date(batch, col::BALANCING_DAY, row)?,
})
}
}
fn pg(e: sqlx::Error) -> Error {
let Some(db) = e.as_database_error() else {
return Error::Storage(e.to_string());
};
match db.code().as_deref() {
Some(code) if code.starts_with("23") => Error::IntegrityViolation {
table: db.table().unwrap_or("<unknown>").to_string(),
constraint: db.constraint().map(str::to_string),
detail: db.message().to_string(),
},
_ => Error::Storage(e.to_string()),
}
}
const LOCK_NOT_AVAILABLE: &str = "55P03";
fn pg_ddl(relation: &str, operation: &str, timeout: Duration, e: sqlx::Error) -> Error {
let timed_out = e
.as_database_error()
.and_then(|db| db.code().map(|c| c == LOCK_NOT_AVAILABLE))
.unwrap_or(false);
match timed_out {
true => Error::LockTimeout {
relation: relation.to_string(),
operation: operation.to_string(),
waited_ms: timeout.whole_milliseconds().max(0) as u64,
},
false => pg(e),
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Attachment {
Any,
Detached,
}
async fn partitions_of(
pool: &PgPool,
table: &str,
attachment: Attachment,
) -> Result<Vec<PartitionId>> {
let sql = format!(
r#"SELECT c.relname
FROM pg_class c
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE c.relkind = 'r'
AND n.nspname = current_schema()
AND c.relname LIKE $1
{}
ORDER BY c.relname"#,
match attachment {
Attachment::Any => "",
Attachment::Detached =>
"AND NOT EXISTS (SELECT 1 FROM pg_inherits i WHERE i.inhrelid = c.oid)",
},
);
let rows = sqlx::query_scalar::<_, String>(&sql)
.bind(format!("{}\\_%", table.replace('_', "\\_")))
.fetch_all(pool)
.await
.map_err(pg)?;
let mut found: Vec<PartitionId> = rows
.iter()
.filter_map(|r| PartitionId::from_relation_name(table, r).ok())
.collect();
found.sort();
Ok(found)
}
#[async_trait]
impl HotStore for PostgresHot {
async fn try_archive_lease(
&self,
table: &str,
) -> Result<Option<Box<dyn crate::tiering::store::TableLease>>> {
let key = lock_key("archive", table);
let mut connection = self.pool.acquire().await.map_err(pg)?;
let acquired = sqlx::query_scalar::<_, bool>("SELECT pg_try_advisory_lock($1)")
.bind(key)
.fetch_one(&mut *connection)
.await
.map_err(pg)?;
if !acquired {
debug!(table, "archive lease held elsewhere");
return Ok(None);
}
debug!(table, "archive lease acquired");
Ok(Some(Box::new(PgTableLease {
connection,
key,
table: table.to_string(),
purpose: "archive",
})))
}
async fn cold_append_lease(
&self,
table: &str,
) -> Result<Box<dyn crate::tiering::store::TableLease>> {
let key = lock_key("cold-append", table);
let mut connection = self.pool.acquire().await.map_err(pg)?;
for attempt in 0..COLD_APPEND_LEASE_ATTEMPTS {
let acquired = sqlx::query_scalar::<_, bool>("SELECT pg_try_advisory_lock($1)")
.bind(key)
.fetch_one(&mut *connection)
.await
.map_err(pg)?;
if acquired {
return Ok(Box::new(PgTableLease {
connection,
key,
table: table.to_string(),
purpose: "cold-append",
}));
}
tokio::time::sleep(std::time::Duration::from_millis(25 << attempt.min(5))).await;
}
Err(Error::LockTimeout {
relation: table.to_string(),
operation: "cold-append claim".to_string(),
waited_ms: COLD_APPEND_LEASE_BUDGET_MS,
})
}
async fn ensure_partitions(
&self,
table: &str,
from: OffsetDateTime,
until: OffsetDateTime,
step: Duration,
) -> Result<Vec<PartitionId>> {
if step <= Duration::ZERO {
return Err(Error::config("partition step must be positive"));
}
let mut created = Vec::new();
let mut start = crate::watermark::align_to_step(from, step)?;
while start < until {
let id = PartitionId::new(table, start);
if self.create_partition(table, &id, step).await? {
created.push(id);
}
start += step;
}
Ok(created)
}
async fn drop_table(&self, table: &str) -> Result<()> {
sqlx::query(&format!(r#"DROP TABLE IF EXISTS "{table}" CASCADE"#))
.execute(&self.pool)
.await
.map_err(pg)?;
for orphan in self.orphaned_partitions(table).await? {
let name = orphan.relation_name()?;
sqlx::query(&format!(r#"DROP TABLE IF EXISTS "{name}""#))
.execute(&self.pool)
.await
.map_err(pg)?;
}
info!(table, "hot table dropped");
Ok(())
}
async fn partition_exists(&self, partition: &PartitionId) -> Result<bool> {
self.relation_exists(&partition.relation_name()?).await
}
async fn detach_partition(&self, partition: &PartitionId) -> Result<()> {
let name = partition.relation_name()?;
let table = partition.table();
let mut tx = self.begin_ddl().await?;
sqlx::query(&format!(
r#"ALTER TABLE "{table}" DETACH PARTITION "{name}""#
))
.execute(&mut *tx)
.await
.map_err(|e| pg_ddl(table, "detach partition", self.ddl_lock_timeout, e))?;
tx.commit().await.map_err(pg)?;
debug!(partition = %name, "detached");
Ok(())
}
async fn create_tables(
&self,
table: &str,
merge_key: &[String],
extra: &[crate::arrow::datatypes::Field],
time_model: crate::config::TimeModel,
) -> Result<()> {
self.create_table_with_key(table, merge_key, extra, time_model)
.await
}
async fn append_reporting(
&self,
table: &str,
merge_key: &[String],
batches: &[RecordBatch],
) -> Result<Vec<crate::session::Displacement>> {
let mut out = Vec::new();
for batch in batches {
out.extend(self.insert_reporting(table, merge_key, batch).await?);
}
Ok(out)
}
async fn append(
&self,
table: &str,
merge_key: &[String],
batches: &[RecordBatch],
) -> Result<u64> {
let mut written = 0u64;
for batch in batches {
if batch.num_rows() == 0 {
continue;
}
written += self.append_batch(table, merge_key, batch).await?;
}
Ok(written)
}
async fn scan_range(
&self,
table: &str,
range: TimeRange,
spec: &ScanSpec,
) -> Result<BatchStream> {
if range.is_empty() {
return Ok(Box::pin(futures::stream::empty()));
}
let partitions = partitions_of(&self.pool, table, Attachment::Any).await?;
let mut streams: Vec<BatchStream> = Vec::new();
for (i, partition) in partitions.iter().enumerate() {
let start = partition.start();
let end = partitions.get(i + 1).map(|next| next.start());
let starts_after_range = range.end().is_some_and(|to| start >= to);
let ends_before_range = match (end, range.start()) {
(Some(end), Some(from)) => end <= from,
_ => false,
};
if starts_after_range || ends_before_range {
continue;
}
let name = partition.relation_name()?;
debug!(table, partition = %name, "scanning a hot partition");
streams.push(self.chunked_scan(&name, range, spec));
}
match streams.len() {
0 => Ok(Box::pin(futures::stream::empty())),
1 => Ok(streams.pop().expect("length checked")),
_ => {
use futures::StreamExt;
Ok(Box::pin(futures::stream::iter(streams).flatten()))
}
}
}
async fn scan_detached(&self, partition: &PartitionId, spec: &ScanSpec) -> Result<BatchStream> {
Ok(self.chunked_scan(&partition.relation_name()?, TimeRange::unbounded(), spec))
}
async fn distinct_malo_ids(&self, partition: &PartitionId) -> Result<Option<u64>> {
let name = partition.relation_name()?;
let count = sqlx::query_scalar::<_, i64>(&format!(
r#"SELECT count(DISTINCT malo_id) FROM "{name}""#
))
.fetch_one(&self.pool)
.await
.map_err(pg)?;
Ok(Some(count.max(0) as u64))
}
async fn drop_partition(&self, partition: &PartitionId) -> Result<()> {
let name = partition.relation_name()?;
let mut tx = self.begin_ddl().await?;
sqlx::query(&format!(r#"DROP TABLE IF EXISTS "{name}""#))
.execute(&mut *tx)
.await
.map_err(|e| pg_ddl(&name, "drop partition", self.ddl_lock_timeout, e))?;
tx.commit().await.map_err(pg)?;
debug!(partition = %name, "dropped");
Ok(())
}
async fn orphaned_partitions(&self, table: &str) -> Result<Vec<PartitionId>> {
partitions_of(&self.pool, table, Attachment::Detached).await
}
async fn partition_starts(&self, table: &str) -> Result<Option<Vec<OffsetDateTime>>> {
Ok(Some(
partitions_of(&self.pool, table, Attachment::Any)
.await?
.into_iter()
.map(|p| p.start())
.collect(),
))
}
async fn invariant_violations(&self, table: &str, watermark: TieringWatermark) -> Result<u64> {
let sql = format!(r#"SELECT count(*) FROM "{table}" WHERE "from" < $1"#);
let count = sqlx::query_scalar::<_, i64>(&sql)
.bind(watermark.get())
.fetch_one(&self.pool)
.await
.map_err(pg)?;
Ok(count as u64)
}
}
fn pg_timestamp(ts: OffsetDateTime) -> Result<String> {
ts.format(&time::format_description::well_known::Rfc3339)
.map_err(|e| Error::encode("partition bound", e.to_string()))
}
fn decimal_to_i128(value: Decimal, scale: i8, column: &str) -> Result<i128> {
let target = u32::try_from(scale).map_err(|_| Error::decode(column, "negative scale"))?;
let mut v = value.normalize();
if v.scale() > target {
return Err(Error::decode(
column,
format!("{value} has more than {scale} decimal places"),
));
}
v.rescale(target);
v.mantissa()
.to_i128()
.ok_or_else(|| Error::decode(column, format!("{value} does not fit i128")))
}
#[cfg(test)]
mod tests {
use super::*;
use time::macros::datetime;
#[test]
fn partitions_are_created_on_the_shared_alignment() {
assert_eq!(
crate::watermark::align_to_step(datetime!(2026-07-20 13:47:03 UTC), Duration::DAY)
.unwrap(),
datetime!(2026-07-20 00:00 UTC)
);
}
#[test]
fn the_cold_append_lease_waits_long_enough_for_a_slow_commit() {
let mut total = 0u64;
for attempt in 0..COLD_APPEND_LEASE_ATTEMPTS {
total += 25u64 << attempt.min(5);
}
assert_eq!(total, COLD_APPEND_LEASE_BUDGET_MS);
assert!(
(10_000..60_000).contains(&COLD_APPEND_LEASE_BUDGET_MS),
"the budget has to cover a slow Iceberg commit under contention and \
still report a genuinely stuck holder rather than wait forever: {}ms",
COLD_APPEND_LEASE_BUDGET_MS
);
let contended = Error::LockTimeout {
relation: "readings_versions".to_string(),
operation: "cold-append claim".to_string(),
waited_ms: COLD_APPEND_LEASE_BUDGET_MS,
};
assert!(contended.is_retryable(), "nothing was changed");
let rendered = contended.to_string();
assert!(rendered.contains("cold-append claim"), "{rendered}");
assert!(rendered.contains("nothing was changed"), "{rendered}");
}
#[test]
fn every_shape_pattern_admits_what_metering_parses_and_no_less() {
use crate::config::{EicType, ValueCheck};
let ok = |c: char| c.is_ascii_digit() || c.is_ascii_uppercase() || c == '-';
let matches = |scheme: ValueCheck, code: &str| {
let b: Vec<char> = code.chars().collect();
match scheme {
ValueCheck::Eic(want) => {
b.len() == 16
&& b.iter().copied().all(ok)
&& b[15] != '-'
&& match want {
Some(ty) => b[2].to_string() == ty.as_str(),
None => b[2].is_ascii_uppercase(),
}
}
ValueCheck::Malo => {
b.len() == 11 && b.iter().all(char::is_ascii_digit) && b[0] != '0'
}
ValueCheck::Melo => {
b.len() == 33
&& b[..2].iter().all(char::is_ascii_uppercase)
&& b[2..8].iter().all(char::is_ascii_digit)
&& b[8..]
.iter()
.all(|c| c.is_ascii_alphanumeric() && !c.is_lowercase())
}
ValueCheck::Bdew => b.len() == 13 && b.iter().all(char::is_ascii_digit),
}
};
assert_eq!(
ValueCheck::ALL.map(ValueCheck::shape_pattern),
[
"^[0-9A-Z-]{2}[A-Z][0-9A-Z-]{12}[0-9A-Z]$",
"^[0-9A-Z-]{2}X[0-9A-Z-]{12}[0-9A-Z]$",
"^[0-9A-Z-]{2}Y[0-9A-Z-]{12}[0-9A-Z]$",
"^[0-9A-Z-]{2}Z[0-9A-Z-]{12}[0-9A-Z]$",
"^[0-9A-Z-]{2}W[0-9A-Z-]{12}[0-9A-Z]$",
"^[0-9A-Z-]{2}T[0-9A-Z-]{12}[0-9A-Z]$",
"^[0-9A-Z-]{2}V[0-9A-Z-]{12}[0-9A-Z]$",
"^[0-9A-Z-]{2}A[0-9A-Z-]{12}[0-9A-Z]$",
"^[1-9][0-9]{10}$",
"^[A-Z]{2}[0-9]{6}[0-9A-Z]{25}$",
"^[0-9]{13}$",
]
);
for scheme in ValueCheck::ALL {
assert_eq!(value_check_pattern(scheme.as_str()), scheme.shape_pattern());
}
for ty in EicType::ALL {
assert_eq!(
ValueCheck::Eic(Some(ty)).shape_pattern(),
ValueCheck::Eic(None)
.shape_pattern()
.replace("[A-Z]", ty.as_str()),
"{ty}"
);
}
let party = ValueCheck::Eic(Some(EicType::Party));
let area = ValueCheck::Eic(Some(EicType::Area));
for (scheme, valid) in [
(ValueCheck::Eic(None), "10X168Y4E6H0041Z"),
(ValueCheck::Eic(None), "10X---ENTSOE---L"),
(ValueCheck::Eic(None), "11XBK0000000001A"),
(ValueCheck::Eic(None), "11YN000000000016"),
(party, "11XBK0000000001A"),
(area, "11YN000000000016"),
(ValueCheck::Malo, "41373559241"),
(ValueCheck::Malo, "12345678905"),
(ValueCheck::Melo, "DE00056266802AO6G56M11SN51G21M24S"),
(ValueCheck::Bdew, "9900987654321"),
(ValueCheck::Bdew, "9812345678901"),
] {
assert_eq!(
scheme.canonicalise("c", valid).ok().as_deref(),
Some(valid),
"{scheme} does not accept {valid}"
);
assert!(matches(scheme, valid), "the DDL would refuse {valid}");
}
for (scheme, bad) in [
(ValueCheck::Eic(None), "11XBK0000000001"), (ValueCheck::Eic(None), "11xbk0000000001a"), (ValueCheck::Eic(None), "11-BK0000000001A"), (ValueCheck::Eic(None), "11XBK0000000001-"),
(party, "11YN000000000016"),
(area, "11XBK0000000001A"),
(ValueCheck::Malo, "0137355924"), (ValueCheck::Malo, "04137355924"), (ValueCheck::Malo, "4137355924A"), (ValueCheck::Melo, "de00056266802AO6G56M11SN51G21M24S"),
(ValueCheck::Melo, "DEX0056266802AO6G56M11SN51G21M24S"), (ValueCheck::Melo, "DE00056266802AO6G56M11SN51G21M24"), (ValueCheck::Bdew, "990098765432"), (ValueCheck::Bdew, "99009876543210"), (ValueCheck::Bdew, "99009876543A1"), ] {
assert!(
!matches(scheme, bad),
"the DDL would accept {bad} as {scheme}"
);
}
}
#[test]
fn an_unknown_value_check_renders_a_pattern_nothing_matches() {
assert_eq!(value_check_pattern("IBAN"), "$^");
}
#[test]
fn decimal_conversion_rejects_excess_precision() {
assert!(decimal_to_i128("0.1".parse().unwrap(), 6, "x").is_ok());
assert!(decimal_to_i128("0.0000001".parse().unwrap(), 6, "x").is_err());
}
#[test]
fn decimal_conversion_scales_correctly() {
assert_eq!(
decimal_to_i128("1.5".parse().unwrap(), 6, "x").unwrap(),
1_500_000
);
assert_eq!(
decimal_to_i128("20260727000001".parse().unwrap(), 0, "v").unwrap(),
20_260_727_000_001
);
}
#[test]
fn a_version_at_the_top_of_the_column_range_round_trips() {
let max = 10i128.pow(20) - 1;
let as_decimal = Decimal::try_from_i128_with_scale(max, 0).unwrap();
assert_eq!(decimal_to_i128(as_decimal, 0, col::VERSION).unwrap(), max);
}
}