aster_forge_test/
postgres.rs1use crate::state::{ContainerLease, ContainerStateLock};
8use crate::suite::TestContainerSuite;
9use sea_orm::{ConnectionTrait, Database, DatabaseConnection};
10use std::time::Duration;
11use testcontainers::core::{ContainerAsync, IntoContainerPort, WaitFor};
12use testcontainers::{GenericImage, ImageExt, ReuseDirective, runners::AsyncRunner};
13
14pub struct PostgresTestContainer {
16 admin_url: String,
17 suite: TestContainerSuite,
18 _container: ContainerAsync<GenericImage>,
19 _lease: ContainerLease,
20}
21
22pub struct PostgresTestDatabase {
24 name: String,
25 url: String,
26 admin_url: String,
27 suite: TestContainerSuite,
28}
29
30impl PostgresTestContainer {
31 pub async fn start(suite: &TestContainerSuite) -> Self {
33 let lock = ContainerStateLock::acquire(suite, "postgres");
34 let mut state = lock.load();
35 let stale_resources = state.prune_stale();
36 state.register_pid(std::process::id());
37 for resource in &stale_resources {
38 state.remember_resource(std::process::id(), resource);
39 }
40 lock.save(&state);
41
42 let container = GenericImage::new("postgres", "16")
43 .with_exposed_port(IntoContainerPort::tcp(5432))
44 .with_wait_for(WaitFor::message_on_stderr(
45 "database system is ready to accept connections",
46 ))
47 .with_container_name(suite.container_name("postgres"))
48 .with_reuse(ReuseDirective::Always)
49 .with_env_var("POSTGRES_USER", "postgres")
50 .with_env_var("POSTGRES_PASSWORD", "postgres")
51 .with_env_var("POSTGRES_DB", "postgres")
52 .start()
53 .await
54 .expect("failed to start PostgreSQL test container");
55 let port = container
56 .get_host_port_ipv4(IntoContainerPort::tcp(5432))
57 .await
58 .expect("PostgreSQL test port should be exposed");
59 drop(lock);
60
61 let fixture = Self {
62 admin_url: format!("postgres://postgres:postgres@127.0.0.1:{port}/postgres"),
63 suite: suite.clone(),
64 _container: container,
65 _lease: ContainerLease::new(suite.clone(), "postgres"),
66 };
67 fixture.cleanup_databases(&stale_resources).await;
68 fixture
69 }
70
71 pub fn admin_url(&self) -> &str {
73 &self.admin_url
74 }
75
76 pub async fn create_database(&self, name: &str) -> PostgresTestDatabase {
78 assert_valid_database_name(name);
79 let lock = ContainerStateLock::acquire(&self.suite, "postgres");
80 let mut state = lock.load();
81 state.remember_resource(std::process::id(), name);
82 lock.save(&state);
83 drop(lock);
84
85 let admin = connect_with_retry(&self.admin_url).await;
86 admin
87 .execute_unprepared(&format!("CREATE DATABASE {}", quote_identifier(name)))
88 .await
89 .unwrap_or_else(|error| {
90 panic!("failed to create PostgreSQL test database {name}: {error}")
91 });
92 admin
93 .close()
94 .await
95 .unwrap_or_else(|error| panic!("failed to close PostgreSQL admin connection: {error}"));
96
97 PostgresTestDatabase {
98 name: name.to_string(),
99 url: database_url(&self.admin_url, name),
100 admin_url: self.admin_url.clone(),
101 suite: self.suite.clone(),
102 }
103 }
104
105 async fn cleanup_databases(&self, names: &[String]) {
106 if names.is_empty() {
107 return;
108 }
109 let admin = connect_with_retry(&self.admin_url).await;
110 for name in names {
111 admin
112 .execute_unprepared(&format!(
113 "DROP DATABASE IF EXISTS {} WITH (FORCE)",
114 quote_identifier(name)
115 ))
116 .await
117 .unwrap_or_else(|error| {
118 panic!("failed to drop stale PostgreSQL test database {name}: {error}")
119 });
120 let lock = ContainerStateLock::acquire(&self.suite, "postgres");
121 let mut state = lock.load();
122 state.forget_resource(std::process::id(), name);
123 lock.save(&state);
124 }
125 admin
126 .close()
127 .await
128 .unwrap_or_else(|error| panic!("failed to close PostgreSQL admin connection: {error}"));
129 }
130}
131
132impl PostgresTestDatabase {
133 pub fn name(&self) -> &str {
135 &self.name
136 }
137
138 pub fn url(&self) -> &str {
140 &self.url
141 }
142
143 pub async fn connect(&self) -> DatabaseConnection {
145 connect_with_retry(&self.url).await
146 }
147
148 pub async fn cleanup(&self) {
150 let admin = connect_with_retry(&self.admin_url).await;
151 admin
152 .execute_unprepared(&format!(
153 "DROP DATABASE IF EXISTS {} WITH (FORCE)",
154 quote_identifier(&self.name)
155 ))
156 .await
157 .unwrap_or_else(|error| {
158 panic!(
159 "failed to drop PostgreSQL test database {}: {error}",
160 self.name
161 )
162 });
163 admin
164 .close()
165 .await
166 .unwrap_or_else(|error| panic!("failed to close PostgreSQL admin connection: {error}"));
167
168 let lock = ContainerStateLock::acquire(&self.suite, "postgres");
169 let mut state = lock.load();
170 state.forget_resource(std::process::id(), &self.name);
171 lock.save(&state);
172 }
173}
174
175async fn connect_with_retry(database_url: &str) -> DatabaseConnection {
176 let deadline = tokio::time::Instant::now() + Duration::from_secs(60);
177 loop {
178 match Database::connect(database_url).await {
179 Ok(database) => return database,
180 Err(error) if tokio::time::Instant::now() >= deadline => {
181 panic!("PostgreSQL test database did not become ready: {error}")
182 }
183 Err(_) => tokio::time::sleep(Duration::from_millis(250)).await,
184 }
185 }
186}
187
188fn database_url(admin_url: &str, name: &str) -> String {
189 admin_url
190 .rsplit_once('/')
191 .map(|(base, _)| format!("{base}/{name}"))
192 .unwrap_or_else(|| admin_url.to_string())
193}
194
195fn quote_identifier(value: &str) -> String {
196 format!("\"{}\"", value.replace('"', "\"\""))
197}
198
199fn assert_valid_database_name(name: &str) {
200 assert!(
201 !name.is_empty()
202 && name.len() <= 63
203 && name
204 .bytes()
205 .all(|byte| byte.is_ascii_alphanumeric() || byte == b'_'),
206 "PostgreSQL test database name must be 1-63 ASCII alphanumeric or '_' characters: {name:?}"
207 );
208}
209
210#[cfg(test)]
211mod tests {
212 use super::{assert_valid_database_name, database_url, quote_identifier};
213
214 #[test]
215 fn database_url_replaces_admin_database() {
216 assert_eq!(
217 database_url("postgres://user:pass@127.0.0.1:5432/postgres", "isolated"),
218 "postgres://user:pass@127.0.0.1:5432/isolated"
219 );
220 }
221
222 #[test]
223 fn identifier_quoting_escapes_quotes() {
224 assert_eq!(quote_identifier("test\"name"), "\"test\"\"name\"");
225 }
226
227 #[test]
228 fn database_name_accepts_boundaries() {
229 assert_valid_database_name("a");
230 assert_valid_database_name(&"a".repeat(63));
231 assert_valid_database_name("aster_product_123");
232 }
233
234 #[test]
235 fn database_name_rejects_unsafe_or_oversized_values() {
236 for name in ["", "has-hyphen", "has quote\"", "has space"] {
237 assert!(
238 std::panic::catch_unwind(|| assert_valid_database_name(name)).is_err(),
239 "database name {name:?} should be rejected"
240 );
241 }
242 let oversized = "a".repeat(64);
243 assert!(std::panic::catch_unwind(|| assert_valid_database_name(&oversized)).is_err());
244 }
245}