docker_wrapper/template/database/
postgres.rs

1//! PostgreSQL template for quick PostgreSQL container setup
2
3#![allow(clippy::doc_markdown)]
4#![allow(clippy::must_use_candidate)]
5#![allow(clippy::return_self_not_must_use)]
6#![allow(clippy::map_unwrap_or)]
7#![allow(clippy::format_push_string)]
8#![allow(clippy::uninlined_format_args)]
9
10use crate::template::{HealthCheck, Template, TemplateConfig, VolumeMount};
11use async_trait::async_trait;
12use std::collections::HashMap;
13
14/// PostgreSQL container template with sensible defaults
15pub struct PostgresTemplate {
16    config: TemplateConfig,
17}
18
19impl PostgresTemplate {
20    /// Create a new PostgreSQL template with default settings
21    pub fn new(name: impl Into<String>) -> Self {
22        let name = name.into();
23        let mut env = HashMap::new();
24
25        // Default PostgreSQL configuration
26        env.insert("POSTGRES_PASSWORD".to_string(), "postgres".to_string());
27        env.insert("POSTGRES_USER".to_string(), "postgres".to_string());
28        env.insert("POSTGRES_DB".to_string(), "postgres".to_string());
29
30        let config = TemplateConfig {
31            name: name.clone(),
32            image: "postgres".to_string(),
33            tag: "15-alpine".to_string(),
34            ports: vec![(5432, 5432)],
35            env,
36            volumes: Vec::new(),
37            network: None,
38            health_check: Some(HealthCheck {
39                test: vec![
40                    "pg_isready".to_string(),
41                    "-U".to_string(),
42                    "postgres".to_string(),
43                ],
44                interval: "10s".to_string(),
45                timeout: "5s".to_string(),
46                retries: 5,
47                start_period: "10s".to_string(),
48            }),
49            auto_remove: false,
50            memory_limit: None,
51            cpu_limit: None,
52            platform: None,
53        };
54
55        Self { config }
56    }
57
58    /// Set a custom PostgreSQL port
59    pub fn port(mut self, port: u16) -> Self {
60        self.config.ports = vec![(port, 5432)];
61        self
62    }
63
64    /// Set database name
65    pub fn database(mut self, db: impl Into<String>) -> Self {
66        self.config.env.insert("POSTGRES_DB".to_string(), db.into());
67        self
68    }
69
70    /// Set database user
71    pub fn user(mut self, user: impl Into<String>) -> Self {
72        let user = user.into();
73        self.config
74            .env
75            .insert("POSTGRES_USER".to_string(), user.clone());
76
77        // Update health check to use the correct user
78        if let Some(health) = &mut self.config.health_check {
79            health.test = vec!["pg_isready".to_string(), "-U".to_string(), user];
80        }
81        self
82    }
83
84    /// Set database password
85    pub fn password(mut self, password: impl Into<String>) -> Self {
86        self.config
87            .env
88            .insert("POSTGRES_PASSWORD".to_string(), password.into());
89        self
90    }
91
92    /// Enable persistence with a volume
93    pub fn with_persistence(mut self, volume_name: impl Into<String>) -> Self {
94        self.config.volumes.push(VolumeMount {
95            source: volume_name.into(),
96            target: "/var/lib/postgresql/data".to_string(),
97            read_only: false,
98        });
99        self
100    }
101
102    /// Mount initialization scripts directory
103    pub fn init_scripts(mut self, scripts_path: impl Into<String>) -> Self {
104        self.config.volumes.push(VolumeMount {
105            source: scripts_path.into(),
106            target: "/docker-entrypoint-initdb.d".to_string(),
107            read_only: true,
108        });
109        self
110    }
111
112    /// Set memory limit for PostgreSQL
113    pub fn memory_limit(mut self, limit: impl Into<String>) -> Self {
114        self.config.memory_limit = Some(limit.into());
115        self
116    }
117
118    /// Set shared memory size
119    pub fn shared_memory(mut self, size: impl Into<String>) -> Self {
120        self.config
121            .env
122            .insert("POSTGRES_SHARED_MEMORY".to_string(), size.into());
123        self
124    }
125
126    /// Enable PostgreSQL extensions
127    pub fn with_extension(mut self, extension: impl Into<String>) -> Self {
128        let ext = extension.into();
129        let current = self
130            .config
131            .env
132            .get("POSTGRES_EXTENSIONS")
133            .map(|s| format!("{},{}", s, ext))
134            .unwrap_or(ext);
135        self.config
136            .env
137            .insert("POSTGRES_EXTENSIONS".to_string(), current);
138        self
139    }
140
141    /// Use a specific PostgreSQL version
142    pub fn version(mut self, version: impl Into<String>) -> Self {
143        self.config.tag = format!("{}-alpine", version.into());
144        self
145    }
146
147    /// Connect to a specific network
148    pub fn network(mut self, network: impl Into<String>) -> Self {
149        self.config.network = Some(network.into());
150        self
151    }
152
153    /// Enable auto-remove when stopped
154    pub fn auto_remove(mut self) -> Self {
155        self.config.auto_remove = true;
156        self
157    }
158
159    /// Set additional PostgreSQL configuration
160    pub fn postgres_args(mut self, args: impl Into<String>) -> Self {
161        self.config
162            .env
163            .insert("POSTGRES_INITDB_ARGS".to_string(), args.into());
164        self
165    }
166
167    /// Enable SSL/TLS
168    pub fn with_ssl(mut self) -> Self {
169        self.config
170            .env
171            .insert("POSTGRES_SSL_MODE".to_string(), "require".to_string());
172        self
173    }
174
175    /// Set locale
176    pub fn locale(mut self, locale: impl Into<String>) -> Self {
177        let locale = locale.into();
178        self.config.env.insert(
179            "POSTGRES_INITDB_ARGS".to_string(),
180            format!("--locale={}", locale),
181        );
182        self
183    }
184
185    /// Use a custom image and tag
186    pub fn custom_image(mut self, image: impl Into<String>, tag: impl Into<String>) -> Self {
187        self.config.image = image.into();
188        self.config.tag = tag.into();
189        self
190    }
191
192    /// Set the platform for the container (e.g., "linux/arm64", "linux/amd64")
193    pub fn platform(mut self, platform: impl Into<String>) -> Self {
194        self.config.platform = Some(platform.into());
195        self
196    }
197}
198
199#[async_trait]
200impl Template for PostgresTemplate {
201    fn name(&self) -> &str {
202        &self.config.name
203    }
204
205    fn config(&self) -> &TemplateConfig {
206        &self.config
207    }
208
209    fn config_mut(&mut self) -> &mut TemplateConfig {
210        &mut self.config
211    }
212
213    async fn wait_for_ready(&self) -> crate::template::Result<()> {
214        use std::time::Duration;
215        use tokio::time::{sleep, timeout};
216
217        // Custom PostgreSQL readiness check
218        // Use 60 second timeout for slower systems (especially Windows)
219        let wait_timeout = Duration::from_secs(60);
220        let check_interval = Duration::from_millis(500);
221
222        timeout(wait_timeout, async {
223            loop {
224                // Check if container is running - keep retrying if not yet started
225                // Don't fail immediately as the container may still be starting up
226                if !self.is_running().await.unwrap_or(false) {
227                    sleep(check_interval).await;
228                    continue;
229                }
230
231                // Try to connect to PostgreSQL using pg_isready
232                let user = self
233                    .config
234                    .env
235                    .get("POSTGRES_USER")
236                    .map(|s| s.as_str())
237                    .unwrap_or("postgres");
238                let db = self
239                    .config
240                    .env
241                    .get("POSTGRES_DB")
242                    .map(|s| s.as_str())
243                    .unwrap_or("postgres");
244
245                let check_cmd = vec!["pg_isready", "-h", "localhost", "-U", user, "-d", db];
246
247                // Execute readiness check
248                if let Ok(result) = self.exec(check_cmd).await {
249                    // pg_isready returns 0 on success
250                    if result.stdout.contains("accepting connections") {
251                        return Ok(());
252                    }
253                }
254
255                sleep(check_interval).await;
256            }
257        })
258        .await
259        .map_err(|_| {
260            crate::template::TemplateError::InvalidConfig(format!(
261                "PostgreSQL container {} failed to become ready within timeout",
262                self.config().name
263            ))
264        })?
265    }
266}
267
268/// Builder for PostgreSQL connection strings
269pub struct PostgresConnectionString {
270    host: String,
271    port: u16,
272    database: String,
273    user: String,
274    password: String,
275}
276
277impl PostgresConnectionString {
278    /// Create from a PostgresTemplate
279    pub fn from_template(template: &PostgresTemplate) -> Self {
280        let config = template.config();
281        let port = config.ports.first().map(|(h, _)| *h).unwrap_or(5432);
282
283        Self {
284            host: "localhost".to_string(),
285            port,
286            database: config
287                .env
288                .get("POSTGRES_DB")
289                .cloned()
290                .unwrap_or_else(|| "postgres".to_string()),
291            user: config
292                .env
293                .get("POSTGRES_USER")
294                .cloned()
295                .unwrap_or_else(|| "postgres".to_string()),
296            password: config
297                .env
298                .get("POSTGRES_PASSWORD")
299                .cloned()
300                .unwrap_or_else(|| "postgres".to_string()),
301        }
302    }
303
304    /// Get the connection string in PostgreSQL URL format
305    pub fn url(&self) -> String {
306        format!(
307            "postgresql://{}:{}@{}:{}/{}",
308            self.user, self.password, self.host, self.port, self.database
309        )
310    }
311
312    /// Get the connection string in key-value format
313    pub fn key_value(&self) -> String {
314        format!(
315            "host={} port={} dbname={} user={} password={}",
316            self.host, self.port, self.database, self.user, self.password
317        )
318    }
319}
320
321#[cfg(test)]
322mod tests {
323    use super::*;
324    use crate::DockerCommand;
325
326    #[test]
327    fn test_postgres_template_basic() {
328        let template = PostgresTemplate::new("test-postgres");
329        assert_eq!(template.name(), "test-postgres");
330        assert_eq!(template.config().image, "postgres");
331        assert_eq!(template.config().tag, "15-alpine");
332        assert_eq!(template.config().ports, vec![(5432, 5432)]);
333    }
334
335    #[test]
336    fn test_postgres_template_custom_config() {
337        let template = PostgresTemplate::new("test-postgres")
338            .database("mydb")
339            .user("myuser")
340            .password("secret123")
341            .port(15432);
342
343        assert_eq!(
344            template.config().env.get("POSTGRES_DB"),
345            Some(&"mydb".to_string())
346        );
347        assert_eq!(
348            template.config().env.get("POSTGRES_USER"),
349            Some(&"myuser".to_string())
350        );
351        assert_eq!(
352            template.config().env.get("POSTGRES_PASSWORD"),
353            Some(&"secret123".to_string())
354        );
355        assert_eq!(template.config().ports, vec![(15432, 5432)]);
356    }
357
358    #[test]
359    fn test_postgres_template_with_persistence() {
360        let template = PostgresTemplate::new("test-postgres").with_persistence("postgres-data");
361
362        assert_eq!(template.config().volumes.len(), 1);
363        assert_eq!(template.config().volumes[0].source, "postgres-data");
364        assert_eq!(
365            template.config().volumes[0].target,
366            "/var/lib/postgresql/data"
367        );
368    }
369
370    #[test]
371    fn test_postgres_template_with_init_scripts() {
372        let template = PostgresTemplate::new("test-postgres").init_scripts("./init-scripts");
373
374        assert_eq!(template.config().volumes.len(), 1);
375        assert_eq!(template.config().volumes[0].source, "./init-scripts");
376        assert_eq!(
377            template.config().volumes[0].target,
378            "/docker-entrypoint-initdb.d"
379        );
380        assert!(template.config().volumes[0].read_only);
381    }
382
383    #[test]
384    fn test_postgres_connection_string() {
385        let template = PostgresTemplate::new("test-postgres")
386            .database("testdb")
387            .user("testuser")
388            .password("testpass")
389            .port(15432);
390
391        let conn = PostgresConnectionString::from_template(&template);
392
393        assert_eq!(
394            conn.url(),
395            "postgresql://testuser:testpass@localhost:15432/testdb"
396        );
397
398        assert_eq!(
399            conn.key_value(),
400            "host=localhost port=15432 dbname=testdb user=testuser password=testpass"
401        );
402    }
403
404    #[test]
405    fn test_postgres_build_command() {
406        let template = PostgresTemplate::new("test-postgres")
407            .database("mydb")
408            .port(15432);
409
410        let cmd = template.build_command();
411        let args = cmd.build_command_args();
412
413        // Check that basic args are present
414        assert!(args.contains(&"run".to_string()));
415        assert!(args.contains(&"--name".to_string()));
416        assert!(args.contains(&"test-postgres".to_string()));
417        assert!(args.contains(&"--publish".to_string()));
418        assert!(args.contains(&"15432:5432".to_string()));
419        assert!(args.contains(&"--env".to_string()));
420    }
421}