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        let wait_timeout = Duration::from_secs(30);
219        let check_interval = Duration::from_millis(500);
220
221        timeout(wait_timeout, async {
222            loop {
223                // First check if container is running
224                if !self.is_running().await? {
225                    return Err(crate::template::TemplateError::NotRunning(
226                        self.config().name.clone(),
227                    ));
228                }
229
230                // Try to connect to PostgreSQL using pg_isready
231                let user = self
232                    .config
233                    .env
234                    .get("POSTGRES_USER")
235                    .map(|s| s.as_str())
236                    .unwrap_or("postgres");
237                let db = self
238                    .config
239                    .env
240                    .get("POSTGRES_DB")
241                    .map(|s| s.as_str())
242                    .unwrap_or("postgres");
243
244                let check_cmd = vec!["pg_isready", "-h", "localhost", "-U", user, "-d", db];
245
246                // Execute readiness check
247                if let Ok(result) = self.exec(check_cmd).await {
248                    // pg_isready returns 0 on success
249                    if result.stdout.contains("accepting connections") {
250                        return Ok(());
251                    }
252                }
253
254                sleep(check_interval).await;
255            }
256        })
257        .await
258        .map_err(|_| {
259            crate::template::TemplateError::InvalidConfig(format!(
260                "PostgreSQL container {} failed to become ready within timeout",
261                self.config().name
262            ))
263        })?
264    }
265}
266
267/// Builder for PostgreSQL connection strings
268pub struct PostgresConnectionString {
269    host: String,
270    port: u16,
271    database: String,
272    user: String,
273    password: String,
274}
275
276impl PostgresConnectionString {
277    /// Create from a PostgresTemplate
278    pub fn from_template(template: &PostgresTemplate) -> Self {
279        let config = template.config();
280        let port = config.ports.first().map(|(h, _)| *h).unwrap_or(5432);
281
282        Self {
283            host: "localhost".to_string(),
284            port,
285            database: config
286                .env
287                .get("POSTGRES_DB")
288                .cloned()
289                .unwrap_or_else(|| "postgres".to_string()),
290            user: config
291                .env
292                .get("POSTGRES_USER")
293                .cloned()
294                .unwrap_or_else(|| "postgres".to_string()),
295            password: config
296                .env
297                .get("POSTGRES_PASSWORD")
298                .cloned()
299                .unwrap_or_else(|| "postgres".to_string()),
300        }
301    }
302
303    /// Get the connection string in PostgreSQL URL format
304    pub fn url(&self) -> String {
305        format!(
306            "postgresql://{}:{}@{}:{}/{}",
307            self.user, self.password, self.host, self.port, self.database
308        )
309    }
310
311    /// Get the connection string in key-value format
312    pub fn key_value(&self) -> String {
313        format!(
314            "host={} port={} dbname={} user={} password={}",
315            self.host, self.port, self.database, self.user, self.password
316        )
317    }
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323    use crate::DockerCommand;
324
325    #[test]
326    fn test_postgres_template_basic() {
327        let template = PostgresTemplate::new("test-postgres");
328        assert_eq!(template.name(), "test-postgres");
329        assert_eq!(template.config().image, "postgres");
330        assert_eq!(template.config().tag, "15-alpine");
331        assert_eq!(template.config().ports, vec![(5432, 5432)]);
332    }
333
334    #[test]
335    fn test_postgres_template_custom_config() {
336        let template = PostgresTemplate::new("test-postgres")
337            .database("mydb")
338            .user("myuser")
339            .password("secret123")
340            .port(15432);
341
342        assert_eq!(
343            template.config().env.get("POSTGRES_DB"),
344            Some(&"mydb".to_string())
345        );
346        assert_eq!(
347            template.config().env.get("POSTGRES_USER"),
348            Some(&"myuser".to_string())
349        );
350        assert_eq!(
351            template.config().env.get("POSTGRES_PASSWORD"),
352            Some(&"secret123".to_string())
353        );
354        assert_eq!(template.config().ports, vec![(15432, 5432)]);
355    }
356
357    #[test]
358    fn test_postgres_template_with_persistence() {
359        let template = PostgresTemplate::new("test-postgres").with_persistence("postgres-data");
360
361        assert_eq!(template.config().volumes.len(), 1);
362        assert_eq!(template.config().volumes[0].source, "postgres-data");
363        assert_eq!(
364            template.config().volumes[0].target,
365            "/var/lib/postgresql/data"
366        );
367    }
368
369    #[test]
370    fn test_postgres_template_with_init_scripts() {
371        let template = PostgresTemplate::new("test-postgres").init_scripts("./init-scripts");
372
373        assert_eq!(template.config().volumes.len(), 1);
374        assert_eq!(template.config().volumes[0].source, "./init-scripts");
375        assert_eq!(
376            template.config().volumes[0].target,
377            "/docker-entrypoint-initdb.d"
378        );
379        assert!(template.config().volumes[0].read_only);
380    }
381
382    #[test]
383    fn test_postgres_connection_string() {
384        let template = PostgresTemplate::new("test-postgres")
385            .database("testdb")
386            .user("testuser")
387            .password("testpass")
388            .port(15432);
389
390        let conn = PostgresConnectionString::from_template(&template);
391
392        assert_eq!(
393            conn.url(),
394            "postgresql://testuser:testpass@localhost:15432/testdb"
395        );
396
397        assert_eq!(
398            conn.key_value(),
399            "host=localhost port=15432 dbname=testdb user=testuser password=testpass"
400        );
401    }
402
403    #[test]
404    fn test_postgres_build_command() {
405        let template = PostgresTemplate::new("test-postgres")
406            .database("mydb")
407            .port(15432);
408
409        let cmd = template.build_command();
410        let args = cmd.build_command_args();
411
412        // Check that basic args are present
413        assert!(args.contains(&"run".to_string()));
414        assert!(args.contains(&"--name".to_string()));
415        assert!(args.contains(&"test-postgres".to_string()));
416        assert!(args.contains(&"--publish".to_string()));
417        assert!(args.contains(&"15432:5432".to_string()));
418        assert!(args.contains(&"--env".to_string()));
419    }
420}