use std::sync::atomic::{AtomicU64, Ordering};
use sqlx::{Connection, PgConnection};
use testcontainers::ImageExt;
use testcontainers::runners::AsyncRunner;
use testcontainers_modules::postgres::Postgres;
use tokio::sync::OnceCell;
use crate::error::{Error, Result};
pub const IMAGE_TAG: &str = "16-alpine";
struct Shared {
admin_url: String,
_container: testcontainers::ContainerAsync<Postgres>,
}
static SHARED: OnceCell<Shared> = OnceCell::const_new();
static NEXT_DATABASE: AtomicU64 = AtomicU64::new(0);
const MAX_CONNECTIONS: &str = "500";
async fn start_container() -> Result<(String, testcontainers::ContainerAsync<Postgres>)> {
let container = Postgres::default()
.with_tag(IMAGE_TAG)
.with_cmd([
"postgres",
"-c",
&format!("max_connections={MAX_CONNECTIONS}"),
])
.start()
.await
.map_err(|e| Error::Storage(format!("starting postgres: {e}")))?;
let port = container
.get_host_port_ipv4(5432)
.await
.map_err(|e| Error::Storage(format!("mapping postgres port: {e}")))?;
Ok((
format!("postgresql://postgres:postgres@127.0.0.1:{port}/postgres"),
container,
))
}
async fn connect_ready(url: &str) -> Result<PgConnection> {
let mut last = None;
for attempt in 0..20u32 {
match PgConnection::connect(url).await {
Ok(connection) => return Ok(connection),
Err(e) => {
last = Some(e);
tokio::time::sleep(std::time::Duration::from_millis(
50 * u64::from(attempt + 1),
))
.await;
}
}
}
Err(Error::Storage(format!(
"postgres did not accept a connection: {}",
last.expect("at least one attempt")
)))
}
pub async fn fresh_database() -> Result<String> {
let shared = SHARED
.get_or_try_init(|| async {
let (admin_url, container) = start_container().await?;
connect_ready(&admin_url).await?.close().await.ok();
Ok::<_, Error>(Shared {
admin_url,
_container: container,
})
})
.await?;
let name = format!(
"meterstore_it_{}",
NEXT_DATABASE.fetch_add(1, Ordering::Relaxed)
);
let mut admin = connect_ready(&shared.admin_url).await?;
sqlx::query(&format!(r#"CREATE DATABASE "{name}""#))
.execute(&mut admin)
.await
.map_err(|e| Error::Storage(format!("creating test database {name}: {e}")))?;
admin.close().await.ok();
Ok(shared
.admin_url
.rsplit_once('/')
.map(|(prefix, _)| format!("{prefix}/{name}"))
.expect("the admin URL carries a database path"))
}
pub struct IsolatedPostgres {
pub url: String,
_container: testcontainers::ContainerAsync<Postgres>,
}
pub async fn isolated_database() -> Result<IsolatedPostgres> {
let (url, container) = start_container().await?;
connect_ready(&url).await?.close().await.ok();
Ok(IsolatedPostgres {
url,
_container: container,
})
}