use std::sync::Arc;
use sqlx::PgPool;
use time::{Duration, OffsetDateTime};
use crate::cold::{IcebergCold, IcebergSqlCatalog, WarehouseAuth};
use crate::config::{TableConfig, ValidatedTableConfig};
use crate::error::Result;
use crate::hot::PostgresHot;
use crate::session::MeterStore;
use crate::tiering::store::{ColdStore, HotStore};
use crate::watermark::ArchivalWindow;
pub struct TestHarness {
hot: Arc<PostgresHot>,
cold: Arc<IcebergCold>,
config: ValidatedTableConfig,
url: String,
_warehouse: tempfile::TempDir,
}
impl std::fmt::Debug for TestHarness {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("TestHarness")
.field("table", &self.config.name())
.finish_non_exhaustive()
}
}
impl TestHarness {
pub const POSTGRES_IMAGE_TAG: &'static str = super::postgres::IMAGE_TAG;
pub const TABLE: &'static str = "readings_versions";
pub async fn start() -> Result<Self> {
Self::with_config(
TableConfig::new(Self::TABLE)
.settlement_lag(Duration::DAY)
.build()?,
)
.await
}
pub async fn with_config(config: ValidatedTableConfig) -> Result<Self> {
let url = super::postgres::fresh_database().await?;
let pool = PgPool::connect(&url)
.await
.map_err(|e| crate::Error::Storage(e.to_string()))?;
let hot = Arc::new(PostgresHot::new(pool));
let warehouse = tempfile::tempdir()
.map_err(|e| crate::Error::Storage(format!("temp warehouse: {e}")))?;
let warehouse_uri = format!("file://{}", warehouse.path().display());
let cold = IcebergSqlCatalog {
database_url: &url,
warehouse_uri: &warehouse_uri,
catalog_name: "meterstore",
namespace: "metering",
file_target_bytes: 8 * 1024 * 1024,
metadata_pool_max_connections: 10,
auth: &WarehouseAuth::default(),
}
.build()
.await?
.cold();
let harness = Self {
hot,
cold,
config,
url,
_warehouse: warehouse,
};
harness
.hot
.create_tables(
harness.config.name(),
&harness.config.merge_key(),
&harness.config.extra_columns(),
harness.config.time_model(),
)
.await?;
harness
.cold
.create_tables(
harness.config.name(),
&harness.config.identity_column_names(),
&harness.config.extra_columns(),
)
.await?;
Ok(harness)
}
pub fn hot(&self) -> &Arc<PostgresHot> {
&self.hot
}
pub fn cold(&self) -> &Arc<IcebergCold> {
&self.cold
}
pub fn config(&self) -> &ValidatedTableConfig {
&self.config
}
pub fn url(&self) -> &str {
&self.url
}
pub fn warehouse(&self) -> &std::path::Path {
self._warehouse.path()
}
pub fn parquet_files(&self) -> Vec<std::path::PathBuf> {
fn walk(dir: &std::path::Path, out: &mut Vec<std::path::PathBuf>) {
let Ok(entries) = std::fs::read_dir(dir) else {
return;
};
for entry in entries.flatten() {
let path = entry.path();
if path.is_dir() {
walk(&path, out);
} else if path.extension().is_some_and(|e| e == "parquet") {
out.push(path);
}
}
}
let mut out = Vec::new();
walk(self._warehouse.path(), &mut out);
out.sort();
out
}
pub async fn store(&self) -> Result<MeterStore> {
self.builder_for(self.config.clone()).await?.build().await
}
pub async fn builder_for(
&self,
config: ValidatedTableConfig,
) -> Result<crate::MeterStoreBuilder> {
self.cold
.create_table_with(
config.name(),
&config.extra_columns(),
&config.identity_column_names(),
)
.await?;
Ok(MeterStore::builder()
.hot(Arc::clone(&self.hot) as Arc<dyn HotStore>)
.cold(
Arc::clone(&self.cold) as Arc<dyn ColdStore>,
self.cold.table_provider(config.name()).await?,
)
.table(config))
}
pub async fn ensure_partitions(&self, from: OffsetDateTime, to: OffsetDateTime) -> Result<()> {
self.hot
.ensure_partitions(self.config.name(), from, to, self.config.archival_step())
.await
.map(|_| ())
}
pub async fn seed_watermark(&self, at: OffsetDateTime) -> Result<()> {
self.seed_watermark_for(self.config.name(), at, self.config.archival_step())
.await
}
pub async fn seed_watermark_for(
&self,
table: &str,
at: OffsetDateTime,
step: time::Duration,
) -> Result<()> {
self.cold
.append_and_commit(
table,
crate::tiering::store::stream_of(Vec::new()),
crate::tiering::store::WriteHints::default(),
ArchivalWindow::new(at - step, at)?,
)
.await
.map(|_| ())
}
pub async fn ingest(
&self,
store: &MeterStore,
series: &[crate::encode::StoredSeries],
) -> Result<()> {
for delivery in series {
store.append(std::slice::from_ref(delivery)).await?;
}
Ok(())
}
pub async fn ingest_readings(
&self,
store: &MeterStore,
deliveries: &[crate::encode::StoredReadings],
) -> Result<()> {
for delivery in deliveries {
store
.append_readings(std::slice::from_ref(delivery))
.await?;
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn the_pinned_image_meets_the_documented_minimum() {
let major: u32 = TestHarness::POSTGRES_IMAGE_TAG
.split('-')
.next()
.and_then(|m| m.parse().ok())
.expect("the tag must start with a major version");
assert!(major >= 12, "the documented floor is PostgreSQL 12");
}
#[test]
fn the_table_name_announces_that_it_holds_versions() {
assert!(TestHarness::TABLE.ends_with("_versions"));
}
}