use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use ring::rand::{SecureRandom, SystemRandom};
use tokio_postgres::{Client, Config};
use crate::backends::control_plane::postgres::{ControlPlaneSettings, PostgresControlPlane};
use crate::backends::control_plane::{ControlPlaneError, ControlPlaneStore};
use crate::desired_state::{
DesiredState, ExpectedRevision, LoadedRevision, RevisionId, TenantId, fixtures,
};
fn role_password() -> String {
let mut bytes = [0u8; 24];
SystemRandom::new()
.fill(&mut bytes)
.expect("the system random generator");
bytes.iter().map(|byte| format!("{byte:02x}")).collect()
}
pub(crate) fn caller() -> TenantId {
fixtures::tenant_id(1)
}
pub(crate) fn other() -> TenantId {
fixtures::tenant_id(11)
}
pub(crate) fn two_tenant_state() -> DesiredState {
let mut state = fixtures::two_tenant_directory_state();
let credential = fixtures::credential(&other(), 13, "secondary");
state
.insert(credential.clone())
.and_then(|state| {
state.insert(fixtures::alias(
&other(),
14,
"steady",
&[credential.reference],
))
})
.and_then(|state| state.insert(fixtures::tenant_policy(1, 1)))
.and_then(|state| state.insert(fixtures::tenant_policy(11, 1)))
.expect("two tenants that reference nothing of each other's are valid");
state
}
pub(crate) fn one_tenant_state() -> DesiredState {
let mut state = fixtures::state_with_directory();
state
.insert(fixtures::tenant_policy(1, 1))
.expect("one tenant's own policy over its own tenant is valid");
state
}
pub(crate) fn two_tenant_catalogue_state() -> DesiredState {
let mut state = two_tenant_state();
let mine = fixtures::tenant_enablement(&caller(), 50, MODEL);
let theirs = fixtures::tenant_enablement(&other(), 60, MODEL);
state
.insert(mine.clone())
.and_then(|state| {
state.insert(fixtures::typed_alias(
&caller(),
&fixtures::project_id(2),
51,
"quick",
&[mine.reference],
))
})
.and_then(|state| state.insert(theirs.clone()))
.and_then(|state| {
state.insert(fixtures::typed_alias(
&other(),
&fixtures::project_id(12),
61,
"swift",
&[theirs.reference],
))
})
.expect("each tenant enabling the same offering for itself is valid");
state
}
pub(crate) const MODEL: &str = "gpt-4o";
pub(crate) struct Journal {
pub(crate) store: Arc<PostgresControlPlane>,
dsn: String,
schema: String,
roles: Mutex<Vec<String>>,
password: String,
}
impl Journal {
pub(crate) async fn open() -> Option<Self> {
let dsn = crate::test_services::postgres_dsn()?;
let schema = format!(
"ti_{}",
SystemTime::now()
.duration_since(UNIX_EPOCH)
.expect("a monotonic wall clock")
.as_nanos()
);
connect(&dsn)
.await
.batch_execute(&format!("CREATE SCHEMA {schema}"))
.await
.expect("a fresh scenario schema");
let store = PostgresControlPlane::connect(
&dsn,
ControlPlaneSettings {
schema: Some(schema.clone()),
operation_timeout: Duration::from_secs(10),
connect_timeout: Duration::from_secs(5),
..ControlPlaneSettings::default()
},
)
.await
.expect("a migrated journal");
Some(Self {
store: Arc::new(store),
dsn,
schema,
roles: Mutex::new(Vec::new()),
password: role_password(),
})
}
pub(crate) fn store(&self) -> Arc<dyn ControlPlaneStore> {
self.store.clone()
}
pub(crate) fn schema(&self) -> &str {
&self.schema
}
pub(crate) async fn publish(
&self,
key: &str,
expected: ExpectedRevision,
state: DesiredState,
) -> Result<RevisionId, ControlPlaneError> {
self.store
.publish_revision(fixtures::candidate(expected, key, state))
.await
.map(|manifest| manifest.id)
}
pub(crate) async fn publish_two_tenants(&self) -> RevisionId {
self.publish("two-tenants", ExpectedRevision::Empty, two_tenant_state())
.await
.expect("two tenants that reference nothing of each other's publish")
}
pub(crate) async fn head(&self) -> Option<RevisionId> {
self.store
.desired_revision()
.await
.expect("the head is readable")
}
pub(crate) async fn hydrated(&self) -> LoadedRevision {
self.store
.load_desired_revision()
.await
.expect("the head hydrates")
.expect("a published head")
}
pub(crate) async fn session(
&self,
label: &str,
privileges: &str,
tenant: Option<TenantId>,
) -> Client {
let role = format!("{}_{label}", self.schema);
let schema = &self.schema;
let password = &self.password;
connect(&self.dsn)
.await
.batch_execute(&format!(
"CREATE ROLE {role} LOGIN PASSWORD '{password}'; \
GRANT USAGE ON SCHEMA {schema} TO {role}; \
GRANT {privileges} ON ALL TABLES IN SCHEMA {schema} TO {role}"
))
.await
.expect("a scenario role");
self.roles.lock().expect("the role list").push(role.clone());
let client = connect_as(&self.dsn, &role, password).await;
let pin = match tenant {
Some(tenant) => format!("SET axond.tenant_id = '{tenant}'"),
None => String::from("RESET axond.tenant_id"),
};
client
.batch_execute(&format!("SET search_path TO {schema}; {pin}"))
.await
.expect("a pinned session");
client
}
pub(crate) async fn stored(&self, sql: &str) -> Vec<String> {
let client = connect(&self.dsn).await;
client
.batch_execute(&format!("SET search_path TO {}", self.schema))
.await
.expect("the journal's schema");
column(&client, sql).await
}
}
impl Drop for Journal {
fn drop(&mut self) {
let dsn = self.dsn.clone();
let schema = self.schema.clone();
let roles = std::mem::take(&mut *self.roles.lock().expect("the role list"));
let cleanup = std::thread::spawn(move || {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.expect("a cleanup runtime");
runtime.block_on(async {
let client = connect(&dsn).await;
let mut left_behind = Vec::new();
for role in roles {
if let Err(error) = client
.batch_execute(&format!(
"DROP OWNED BY {role} CASCADE; DROP ROLE IF EXISTS {role}"
))
.await
{
left_behind.push(format!("the login role {role}: {}", detail(&error)));
}
}
if let Err(error) = client
.batch_execute(&format!("DROP SCHEMA IF EXISTS {schema} CASCADE"))
.await
{
left_behind.push(format!("the schema {schema}: {}", detail(&error)));
}
assert!(
left_behind.is_empty(),
"a scenario could not clean up after itself: {left_behind:?}"
);
});
});
if cleanup.join().is_err() && !std::thread::panicking() {
panic!("a scenario left its schema or its login roles behind");
}
}
}
async fn connect(dsn: &str) -> Client {
open(dsn.parse().expect("a parseable test DSN")).await
}
async fn connect_as(dsn: &str, role: &str, password: &str) -> Client {
let mut config: Config = dsn.parse().expect("a parseable test DSN");
config.user(role).password(password);
open(config).await
}
async fn open(mut config: Config) -> Client {
config.connect_timeout(Duration::from_secs(5));
let (client, connection) = config
.connect(crate::usage::tls_connector())
.await
.expect("a connection to the test database");
tokio::spawn(async move {
let _ = connection.await;
});
client
}
pub(crate) fn detail(error: &tokio_postgres::Error) -> String {
let mut rendered = error.to_string();
let mut source = std::error::Error::source(error);
while let Some(cause) = source {
rendered.push_str(": ");
rendered.push_str(&cause.to_string());
source = cause.source();
}
rendered
}
pub(crate) async fn column(client: &Client, sql: &str) -> Vec<String> {
client
.query(sql, &[])
.await
.unwrap_or_else(|error| panic!("the read itself must succeed: {}", detail(&error)))
.iter()
.map(|row| {
row.try_get::<_, Option<String>>(0)
.expect("a text column")
.unwrap_or_else(|| "<null>".to_owned())
})
.collect()
}
pub(crate) async fn refused(client: &Client, sql: &str) -> String {
let refusal = client
.batch_execute(sql)
.await
.expect_err("the database must refuse the write");
detail(&refusal)
}
pub(crate) async fn affected(client: &Client, sql: &str) -> u64 {
client
.execute(sql, &[])
.await
.unwrap_or_else(|error| panic!("the statement itself must run: {}", detail(&error)))
}
pub(crate) struct Absent {
names: Vec<(&'static str, String)>,
}
impl Absent {
pub(crate) fn of(names: impl IntoIterator<Item = (&'static str, String)>) -> Self {
let names: Vec<_> = names.into_iter().collect();
assert!(
names.iter().all(|(_, value)| !value.is_empty()),
"an empty identifier would make every absence assertion vacuous"
);
Self { names }
}
pub(crate) fn of_the_other_tenant() -> Self {
let credential = fixtures::credential(&other(), 13, "secondary");
Self::of([
("tenant id", other().to_string()),
("tenant slug", "globex".to_owned()),
("project id", fixtures::project_id(12).to_string()),
("principal id", fixtures::principal_id(40).to_string()),
("workload id", fixtures::principal_id(41).to_string()),
("credential id", credential.reference.id.to_string()),
("secret id", fixtures::secret_id(13).to_string()),
])
}
pub(crate) fn of_the_other_tenants_own_rows() -> Self {
let all = Self::of_the_other_tenant();
Self {
names: all
.names
.into_iter()
.filter(|(label, _)| !matches!(*label, "tenant id" | "tenant slug"))
.collect(),
}
}
pub(crate) fn names(&self) -> &[(&'static str, String)] {
&self.names
}
pub(crate) fn assert_absent(&self, surface: &str, rendered: &str) {
for (label, value) in &self.names {
assert!(
!rendered.contains(value.as_str()),
"{surface} discloses the other tenant's {label}: {rendered}"
);
}
}
}