1#![deny(missing_docs)]
11#![warn(clippy::all, clippy::pedantic)]
12
13#[cfg(feature = "tokio-process")]
15pub mod asynchronous;
16pub mod errors;
18mod search;
19pub mod synchronous;
21
22use std::fmt::Write as FmtWrite;
23use std::fs::{metadata, set_permissions};
24use std::io::{BufRead, BufReader};
25use std::path::Path;
26use std::sync::atomic::AtomicU32;
27use std::sync::{Arc, OnceLock, RwLock};
28use std::{fs::File, io::Write};
29
30use ctor::dtor;
31use nix::unistd::{Gid, Uid};
32use tempfile::{Builder, TempDir};
33use tracing::{debug, info, instrument};
34
35use crate::errors::{TmpPostgrustError, TmpPostgrustResult};
36
37const TMP_POSTGRUST_DB_NAME: &str = "tmp-postgrust";
38const TMP_POSTGRUST_USER_NAME: &str = "tmp-postgrust-user";
39
40pub(crate) static POSTGRES_UID_GID: OnceLock<(Uid, Gid)> = OnceLock::new();
41
42#[dtor]
46fn cleanup_static() {
47 #[cfg(feature = "tokio-process")]
48 if let Some(factory_mutex) = TOKIO_POSTGRES_FACTORY.get() {
49 let mut guard = factory_mutex.blocking_write();
50 drop(guard.take());
51 }
52
53 if let Some(factory_mutex) = DEFAULT_POSTGRES_FACTORY.get() {
54 let mut guard = factory_mutex
55 .write()
56 .expect("Failed to lock default factory mutex.");
57 drop(guard.take());
58 }
59}
60
61static DEFAULT_POSTGRES_FACTORY: OnceLock<RwLock<Option<TmpPostgrustFactory>>> = OnceLock::new();
62
63pub fn new_default_process() -> TmpPostgrustResult<synchronous::ProcessGuard> {
77 let factory_mutex = DEFAULT_POSTGRES_FACTORY.get_or_init(|| {
78 RwLock::new(Some(
79 TmpPostgrustFactory::try_new(&TmpPostgrustFactoryConfig {
80 disable_fsync: true,
81 })
82 .expect("Failed to initialize default postgres factory."),
83 ))
84 });
85 let guard = factory_mutex
86 .read()
87 .expect("Failed to lock default factory mutex.");
88 let factory = guard
89 .as_ref()
90 .expect("Default factory is uninitialized or has been dropped.");
91 factory.new_instance()
92}
93
94pub fn new_default_process_with_migrations(
109 migrate: impl Fn(&str) -> Result<(), Box<dyn std::error::Error + Send + Sync>>,
110) -> TmpPostgrustResult<synchronous::ProcessGuard> {
111 let factory_mutex = DEFAULT_POSTGRES_FACTORY.get_or_init(|| {
112 let factory = TmpPostgrustFactory::try_new(&TmpPostgrustFactoryConfig {
113 disable_fsync: true,
114 })
115 .expect("Failed to initialize default postgres factory.");
116 factory
117 .run_migrations(migrate)
118 .expect("Failed to run migrations");
119
120 RwLock::new(Some(factory))
121 });
122 let guard = factory_mutex
123 .read()
124 .expect("Failed to lock default factory mutex.");
125 let factory = guard
126 .as_ref()
127 .expect("Default factory is uninitialized or has been dropped.");
128 factory.new_instance()
129}
130
131#[cfg(feature = "tokio-process")]
133static TOKIO_POSTGRES_FACTORY: tokio::sync::OnceCell<
134 tokio::sync::RwLock<Option<TmpPostgrustFactory>>,
135> = tokio::sync::OnceCell::const_new();
136
137#[cfg(feature = "tokio-process")]
151pub async fn new_default_process_async() -> TmpPostgrustResult<asynchronous::ProcessGuard> {
152 let factory_mutex = TOKIO_POSTGRES_FACTORY
153 .get_or_try_init(|| async {
154 TmpPostgrustFactory::try_new_async(&TmpPostgrustFactoryConfig {
155 disable_fsync: true,
156 })
157 .await
158 .map(|factory| tokio::sync::RwLock::new(Some(factory)))
159 })
160 .await?;
161 let guard = factory_mutex.read().await;
162 let factory = guard
163 .as_ref()
164 .expect("Default tokio factory is uninitialized or has been dropped.");
165 factory.new_instance_async().await
166}
167
168#[cfg(feature = "tokio-process")]
183pub async fn new_default_process_async_with_migrations<F>(
184 migrate: F,
185) -> TmpPostgrustResult<asynchronous::ProcessGuard>
186where
187 F: for<'r> Fn(
188 &'r str,
189 ) -> futures_util::future::BoxFuture<
190 'r,
191 Result<(), Box<dyn std::error::Error + Send + Sync>>,
192 >,
193{
194 let factory_mutex = TOKIO_POSTGRES_FACTORY
195 .get_or_try_init(|| async {
196 let factory = TmpPostgrustFactory::try_new_async(&TmpPostgrustFactoryConfig {
197 disable_fsync: true,
198 })
199 .await?;
200 factory.run_migrations_async(migrate).await?;
201 Ok(tokio::sync::RwLock::new(Some(factory)))
202 })
203 .await?;
204 let guard = factory_mutex.read().await;
205 let factory = guard
206 .as_ref()
207 .expect("Default tokio factory is uninitialized or has been dropped.");
208 factory.new_instance_async().await
209}
210
211#[derive(Debug)]
213pub struct TmpPostgrustFactory {
214 socket_dir: Arc<TempDir>,
215 cache_dir: Arc<TempDir>,
216 config: String,
217 next_port: AtomicU32,
218}
219
220#[derive(Default, Debug)]
222pub struct TmpPostgrustFactoryConfig {
223 pub disable_fsync: bool,
227}
228
229impl TmpPostgrustFactory {
230 fn build_config(factory_config: &TmpPostgrustFactoryConfig, socket_dir: &Path) -> String {
232 let mut config = String::new();
233 config.push_str("shared_buffers = '12MB'\n");
235 config.push_str("listen_addresses = ''\n");
237 writeln!(
239 &mut config,
240 "unix_socket_directories = \'{}\'",
241 socket_dir.to_str().unwrap()
242 )
243 .expect("Failed to append unix socket to config");
244 writeln!(&mut config, "fsync = \'{}\'", !factory_config.disable_fsync)
248 .expect("Failed to append fsync option to config");
249
250 config
251 }
252
253 #[instrument]
255 pub fn try_new(
256 factory_config: &TmpPostgrustFactoryConfig,
257 ) -> TmpPostgrustResult<TmpPostgrustFactory> {
258 let socket_dir = Builder::new()
259 .prefix("tmp-postgrust-socket")
260 .tempdir()
261 .map_err(TmpPostgrustError::CreateSocketDirFailed)?;
262 let cache_dir = Builder::new()
263 .prefix("tmp-postgrust-cache")
264 .tempdir()
265 .map_err(TmpPostgrustError::CreateCacheDirFailed)?;
266
267 synchronous::chown_to_non_root(cache_dir.path())?;
268 synchronous::chown_to_non_root(socket_dir.path())?;
269 synchronous::exec_init_db(cache_dir.path())?;
270
271 let config = TmpPostgrustFactory::build_config(factory_config, socket_dir.path());
272
273 let factory = TmpPostgrustFactory {
274 socket_dir: Arc::new(socket_dir),
275 cache_dir: Arc::new(cache_dir),
276 config,
277 next_port: AtomicU32::new(5432),
278 };
279 let process = factory.start_postgresql(&factory.cache_dir)?;
280 synchronous::exec_create_user(process.socket_dir.path(), process.port, &process.user_name)?;
281 synchronous::exec_create_db(
282 process.socket_dir.path(),
283 process.port,
284 &process.user_name,
285 &process.db_name,
286 )?;
287
288 Ok(factory)
289 }
290
291 #[cfg(feature = "tokio-process")]
293 #[instrument]
294 pub async fn try_new_async(
295 factory_config: &TmpPostgrustFactoryConfig,
296 ) -> TmpPostgrustResult<TmpPostgrustFactory> {
297 let socket_dir = Builder::new()
298 .prefix("tmp-postgrust-socket")
299 .tempdir()
300 .map_err(TmpPostgrustError::CreateSocketDirFailed)?;
301 let cache_dir = Builder::new()
302 .prefix("tmp-postgrust-cache")
303 .tempdir()
304 .map_err(TmpPostgrustError::CreateCacheDirFailed)?;
305
306 asynchronous::chown_to_non_root(cache_dir.path()).await?;
307 asynchronous::chown_to_non_root(socket_dir.path()).await?;
308 asynchronous::exec_init_db(cache_dir.path()).await?;
309
310 let config = TmpPostgrustFactory::build_config(factory_config, socket_dir.path());
311
312 let factory = TmpPostgrustFactory {
313 socket_dir: Arc::new(socket_dir),
314 cache_dir: Arc::new(cache_dir),
315 config,
316 next_port: AtomicU32::new(5432),
317 };
318 let process = factory.start_postgresql_async(&factory.cache_dir).await?;
319 asynchronous::exec_create_user(process.socket_dir.path(), process.port, &process.user_name)
320 .await?;
321 asynchronous::exec_create_db(
322 process.socket_dir.path(),
323 process.port,
324 &process.user_name,
325 &process.db_name,
326 )
327 .await?;
328
329 Ok(factory)
330 }
331
332 pub fn run_migrations(
340 &self,
341 migrate: impl Fn(&str) -> Result<(), Box<dyn std::error::Error + Send + Sync>>,
342 ) -> TmpPostgrustResult<()> {
343 let process = self.start_postgresql(&self.cache_dir)?;
344
345 migrate(&process.connection_string()).map_err(TmpPostgrustError::MigrationsFailed)?;
346
347 Ok(())
348 }
349
350 #[cfg(feature = "tokio-process")]
358 pub async fn run_migrations_async<F>(&self, migrate: F) -> TmpPostgrustResult<()>
359 where
360 F: for<'r> Fn(
361 &'r str,
362 ) -> futures_util::future::BoxFuture<
363 'r,
364 Result<(), Box<dyn std::error::Error + Send + Sync>>,
365 >,
366 {
367 let process = self.start_postgresql(&self.cache_dir)?;
368
369 migrate(&process.connection_string())
370 .await
371 .map_err(TmpPostgrustError::MigrationsFailed)?;
372
373 Ok(())
374 }
375
376 #[instrument(skip(self))]
379 pub fn new_instance(&self) -> TmpPostgrustResult<synchronous::ProcessGuard> {
380 let data_directory = Builder::new()
381 .prefix("tmp-postgrust-db")
382 .tempdir()
383 .map_err(TmpPostgrustError::CreateDataDirFailed)?;
384 let data_directory_path = data_directory.path();
385
386 set_permissions(
387 &data_directory,
388 metadata(self.cache_dir.path()).unwrap().permissions(),
389 )
390 .unwrap();
391 synchronous::exec_copy_dir(self.cache_dir.path(), data_directory_path)?;
392
393 if !data_directory_path.join("PG_VERSION").exists() {
394 return Err(TmpPostgrustError::EmptyDataDirectory);
395 }
396
397 self.start_postgresql(&Arc::new(data_directory))
398 }
399
400 #[instrument(skip(self))]
401 fn start_postgresql(
402 &self,
403 dir: &Arc<TempDir>,
404 ) -> TmpPostgrustResult<synchronous::ProcessGuard> {
405 File::create(dir.path().join("postgresql.conf"))
406 .map_err(TmpPostgrustError::CreateConfigFailed)?
407 .write_all(self.config.as_bytes())
408 .map_err(TmpPostgrustError::CreateConfigFailed)?;
409
410 let port = self
411 .next_port
412 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
413
414 synchronous::chown_to_non_root(dir.path())?;
415 let mut postgres_process_handle = synchronous::start_postgres_subprocess(dir.path(), port)?;
416 let stdout = postgres_process_handle.stdout.take().unwrap();
417 let stderr = postgres_process_handle.stderr.take().unwrap();
418
419 let stdout_reader = BufReader::new(stdout).lines();
420 let mut stderr_reader = BufReader::new(stderr).lines();
421
422 while let Some(Ok(line)) = stderr_reader.next() {
423 debug!("Postgresql: {}", line);
424 if line.contains("database system is ready to accept connections") {
425 info!("temporary database system is read to accept connections");
426 break;
427 }
428 }
429
430 Ok(synchronous::ProcessGuard {
431 stdout_reader: Some(stdout_reader),
432 stderr_reader: Some(stderr_reader),
433 port,
434 db_name: TMP_POSTGRUST_DB_NAME.to_string(),
435 user_name: TMP_POSTGRUST_USER_NAME.to_string(),
436 postgres_process: postgres_process_handle,
437 _data_directory: Arc::clone(dir),
438 _cache_directory: Arc::clone(&self.cache_dir),
439 socket_dir: Arc::clone(&self.socket_dir),
440 })
441 }
442
443 #[cfg(feature = "tokio-process")]
446 #[instrument(skip(self))]
447 pub async fn new_instance_async(&self) -> TmpPostgrustResult<asynchronous::ProcessGuard> {
448 use tokio::fs::{metadata, set_permissions};
449
450 let data_directory = Builder::new()
451 .prefix("tmp-postgrust-db")
452 .tempdir()
453 .map_err(TmpPostgrustError::CreateDataDirFailed)?;
454 let data_directory_path = data_directory.path();
455
456 set_permissions(
457 &data_directory,
458 metadata(self.cache_dir.path()).await.unwrap().permissions(),
459 )
460 .await
461 .unwrap();
462 asynchronous::exec_copy_dir(self.cache_dir.path(), data_directory_path).await?;
463
464 if !data_directory_path.join("PG_VERSION").exists() {
465 return Err(TmpPostgrustError::EmptyDataDirectory);
466 }
467
468 self.start_postgresql_async(&Arc::new(data_directory)).await
469 }
470
471 #[cfg(feature = "tokio-process")]
472 #[instrument(skip(self))]
473 async fn start_postgresql_async(
474 &self,
475 dir: &Arc<TempDir>,
476 ) -> TmpPostgrustResult<asynchronous::ProcessGuard> {
477 use tokio::io::AsyncBufReadExt;
478
479 File::create(dir.path().join("postgresql.conf"))
480 .map_err(TmpPostgrustError::CreateConfigFailed)?
481 .write_all(self.config.as_bytes())
482 .map_err(TmpPostgrustError::CreateConfigFailed)?;
483
484 let port = self
485 .next_port
486 .fetch_add(1, std::sync::atomic::Ordering::SeqCst);
487
488 asynchronous::chown_to_non_root(dir.path()).await?;
489 let mut postgres_process_handle =
490 asynchronous::start_postgres_subprocess(dir.path(), port)?;
491 let stdout = postgres_process_handle.stdout.take().unwrap();
492 let stderr = postgres_process_handle.stderr.take().unwrap();
493
494 let stdout_reader = tokio::io::BufReader::new(stdout).lines();
495 let mut stderr_reader = tokio::io::BufReader::new(stderr).lines();
496
497 while let Some(line) = stderr_reader.next_line().await.unwrap() {
498 debug!("Postgresql: {}", line);
499 if line.contains("database system is ready to accept connections") {
500 info!("temporary database system is read to accept connections");
501 break;
502 }
503 }
504
505 Ok(asynchronous::ProcessGuard {
506 stdout_reader: Some(stdout_reader),
507 stderr_reader: Some(stderr_reader),
508 port,
509 db_name: TMP_POSTGRUST_DB_NAME.to_string(),
510 user_name: TMP_POSTGRUST_USER_NAME.to_string(),
511 _data_directory: Arc::clone(dir),
512 _cache_directory: Arc::clone(&self.cache_dir),
513 socket_dir: Arc::clone(&self.socket_dir),
514 postgres_process: postgres_process_handle,
515 })
516 }
517}
518
519#[cfg(test)]
520mod tests {
521 use super::*;
522
523 use test_log::test;
524 use tokio_postgres::NoTls;
525 use tracing::error;
526
527 #[test(tokio::test)]
528 async fn it_works() {
529 let factory = TmpPostgrustFactory::try_new(&TmpPostgrustFactoryConfig {
530 disable_fsync: false,
531 })
532 .expect("failed to create factory");
533
534 let postgresql_proc = factory
535 .new_instance()
536 .expect("failed to create a new instance");
537
538 let (client, conn) = tokio_postgres::connect(&postgresql_proc.connection_string(), NoTls)
539 .await
540 .expect("failed to connect to postgresql");
541
542 tokio::spawn(async move {
543 if let Err(e) = conn.await {
544 error!("connection error: {}", e);
545 }
546 });
547
548 client.query("SELECT 1;", &[]).await.unwrap();
549 }
550
551 #[test(tokio::test)]
552 async fn it_works_fsync_disabled() {
553 let factory = TmpPostgrustFactory::try_new(&TmpPostgrustFactoryConfig {
554 disable_fsync: true,
555 })
556 .expect("failed to create factory");
557
558 let postgresql_proc = factory
559 .new_instance()
560 .expect("failed to create a new instance");
561
562 let (client, conn) = tokio_postgres::connect(&postgresql_proc.connection_string(), NoTls)
563 .await
564 .expect("failed to connect to postgresql");
565
566 tokio::spawn(async move {
567 if let Err(e) = conn.await {
568 error!("connection error: {}", e);
569 }
570 });
571
572 client.query("SELECT 1;", &[]).await.unwrap();
573 }
574
575 #[cfg(feature = "tokio-process")]
576 #[test(tokio::test)]
577 async fn it_works_async() {
578 let factory = TmpPostgrustFactory::try_new_async(&TmpPostgrustFactoryConfig {
579 disable_fsync: false,
580 })
581 .await
582 .expect("failed to create factory");
583
584 let postgresql_proc = factory
585 .new_instance_async()
586 .await
587 .expect("failed to create a new instance");
588
589 let (client, conn) = tokio_postgres::connect(&postgresql_proc.connection_string(), NoTls)
590 .await
591 .expect("failed to connect to postgresql");
592
593 tokio::spawn(async move {
594 if let Err(e) = conn.await {
595 error!("connection error: {}", e);
596 }
597 });
598
599 client.query("SELECT 1;", &[]).await.unwrap();
600 }
601
602 #[test(tokio::test)]
603 async fn two_simulatenous_processes() {
604 let factory = TmpPostgrustFactory::try_new(&TmpPostgrustFactoryConfig {
605 disable_fsync: false,
606 })
607 .expect("failed to create factory");
608
609 let proc1 = factory
610 .new_instance()
611 .expect("failed to create a new instance");
612
613 let proc2 = factory
614 .new_instance()
615 .expect("failed to create a new instance");
616
617 let (client1, conn1) = tokio_postgres::connect(&proc1.connection_string(), NoTls)
618 .await
619 .expect("failed to connect to postgresql");
620
621 tokio::spawn(async move {
622 if let Err(e) = conn1.await {
623 error!("connection error: {}", e);
624 }
625 });
626
627 let (client2, conn2) = tokio_postgres::connect(&proc2.connection_string(), NoTls)
628 .await
629 .expect("failed to connect to postgresql");
630
631 tokio::spawn(async move {
632 if let Err(e) = conn2.await {
633 error!("connection error: {}", e);
634 }
635 });
636
637 client1.query("SELECT 1;", &[]).await.unwrap();
638 client2.query("SELECT 1;", &[]).await.unwrap();
639 }
640
641 #[cfg(feature = "tokio-process")]
642 #[test(tokio::test)]
643 async fn two_simulatenous_processes_async() {
644 let factory = TmpPostgrustFactory::try_new_async(&TmpPostgrustFactoryConfig {
645 disable_fsync: false,
646 })
647 .await
648 .expect("failed to create factory");
649
650 let proc1 = factory
651 .new_instance_async()
652 .await
653 .expect("failed to create a new instance");
654
655 let proc2 = factory
656 .new_instance_async()
657 .await
658 .expect("failed to create a new instance");
659
660 let (client1, conn1) = tokio_postgres::connect(&proc1.connection_string(), NoTls)
661 .await
662 .expect("failed to connect to postgresql");
663
664 tokio::spawn(async move {
665 if let Err(e) = conn1.await {
666 error!("connection error: {}", e);
667 }
668 });
669
670 let (client2, conn2) = tokio_postgres::connect(&proc2.connection_string(), NoTls)
671 .await
672 .expect("failed to connect to postgresql");
673
674 tokio::spawn(async move {
675 if let Err(e) = conn2.await {
676 error!("connection error: {}", e);
677 }
678 });
679
680 client1.query("SELECT 1;", &[]).await.unwrap();
681 client2.query("SELECT 1;", &[]).await.unwrap();
682 }
683
684 #[cfg(feature = "tokio-process")]
685 #[test(tokio::test)]
686 async fn default_process_factory_1() {
687 let proc = new_default_process_async().await.unwrap();
688
689 let (client, conn) = tokio_postgres::connect(&proc.connection_string(), NoTls)
690 .await
691 .expect("failed to connect to postgresql");
692
693 tokio::spawn(async move {
694 if let Err(e) = conn.await {
695 error!("connection error: {}", e);
696 }
697 });
698
699 client.execute("CREATE TABLE lock ();", &[]).await.unwrap();
701 }
702
703 #[cfg(feature = "tokio-process")]
704 #[test(tokio::test)]
705 async fn default_process_factory_2() {
706 let proc = new_default_process_async().await.unwrap();
707
708 let (client, conn) = tokio_postgres::connect(&proc.connection_string(), NoTls)
709 .await
710 .expect("failed to connect to postgresql");
711
712 tokio::spawn(async move {
713 if let Err(e) = conn.await {
714 error!("connection error: {}", e);
715 }
716 });
717
718 client.execute("CREATE TABLE lock ();", &[]).await.unwrap();
720 }
721
722 #[cfg(feature = "tokio-process")]
723 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
724 async fn default_process_factory_multithread_1() {
725 let proc = new_default_process_async().await.unwrap();
726
727 let (client, conn) = tokio_postgres::connect(&proc.connection_string(), NoTls)
728 .await
729 .expect("failed to connect to postgresql");
730
731 tokio::spawn(async move {
732 if let Err(e) = conn.await {
733 error!("connection error: {}", e);
734 }
735 });
736
737 client.execute("CREATE TABLE lock ();", &[]).await.unwrap();
739 }
740
741 #[cfg(feature = "tokio-process")]
742 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
743 async fn default_process_factory_multithread_2() {
744 let proc = new_default_process_async().await.unwrap();
745
746 let (client, conn) = tokio_postgres::connect(&proc.connection_string(), NoTls)
747 .await
748 .expect("failed to connect to postgresql");
749
750 tokio::spawn(async move {
751 if let Err(e) = conn.await {
752 error!("connection error: {}", e);
753 }
754 });
755
756 client.execute("CREATE TABLE lock ();", &[]).await.unwrap();
758 }
759}