temp_mongo/
temp_mongo.rs

1use std::ffi::{OsStr, OsString};
2use std::path::{Path, PathBuf};
3use std::time::Duration;
4
5use crate::error::ErrorInner;
6use crate::util::{KillOnDrop, TempDir};
7use crate::Error;
8use std::process::Command;
9
10/// A temporary MongoDB instance.
11///
12/// All state of the MongoDB instance is stored in a temporary directory.
13/// Unless disabled, the temporary directory is deleted when this object is dropped.
14pub struct TempMongo {
15	tempdir: TempDir,
16	socket_path: PathBuf,
17	log_path: PathBuf,
18	client: mongodb::Client,
19	server: KillOnDrop,
20}
21
22impl std::fmt::Debug for TempMongo {
23	fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
24		f.debug_struct("TempMongo")
25			.field("tempdir", &self.tempdir.path())
26			.field("socket_path", &self.socket_path())
27			.field("log_path", &self.log_path())
28			.field("server_pid", &self.server.id())
29			.finish_non_exhaustive()
30	}
31}
32
33impl TempMongo {
34	/// Spawn a new MongoDB instance with a temporary state directory.
35	pub async fn new() -> Result<Self, Error> {
36		Self::from_builder(&TempMongoBuilder::new()).await
37	}
38
39	/// Create a builder to customize your [`TempMongo`].
40	///
41	/// After configuring the desirec options, run [`TempMongoBuilder::spawn()`].
42	pub fn builder() -> TempMongoBuilder {
43		TempMongoBuilder::new()
44	}
45
46	/// Get the PID of the MongoDB process.
47	pub fn process_id(&self) -> u32 {
48		self.server.id()
49	}
50
51	/// Get the path of the temporary state directory.
52	pub fn directory(&self) -> &Path {
53		self.tempdir.path()
54	}
55
56	/// Get the path of the listening socket of the MongoDB instance.
57	pub fn socket_path(&self) -> &Path {
58		&self.socket_path
59	}
60
61	/// Get the path of the log file of the MongoDB instance.
62	pub fn log_path(&self) -> &Path {
63		&self.log_path
64	}
65
66	/// Get a client for the MongDB instance.
67	///
68	/// This returns a client by reference,
69	/// but it can be cloned and sent to other threads or tasks if needed.
70	pub fn client(&self) -> &mongodb::Client {
71		&self.client
72	}
73
74	/// Enable or disable clean-up of the temporary directory when this object is dropped.
75	pub fn set_clean_on_drop(&mut self, clean_on_drop: bool) {
76		self.tempdir.set_clean_on_drop(clean_on_drop);
77	}
78
79	/// Kill the server and remove the temporary state directory on the filesystem.
80	///
81	/// Note that the server will also be killed when this object is dropped,
82	/// and unless disabled, the temporary state directory will be removed by the [`Drop`] implementation too.
83	///
84	/// This function ignores the value of `clean_on_drop`.
85	/// It also allows for better error handling compared to just dropping the object.
86	pub async fn kill_and_clean(mut self) -> Result<(), Error> {
87		self.client.shutdown_immediate().await;
88		self.server.kill()
89			.map_err(ErrorInner::KillServer)?;
90
91		let path = self.tempdir.path().to_owned();
92		self.tempdir.close()
93			.map_err(|e| ErrorInner::CleanDir(path, e))?;
94		Ok(())
95	}
96
97	/// Kill the server, but leave the temporary state directory on the filesystem.
98	///
99	/// Note that the server will also be killed when this object is dropped.
100	///
101	/// This function ignores the value of `clean_on_drop`.
102	/// It also allows for better error handling compared to just dropping the object.
103	pub async fn kill_no_clean(mut self) -> Result<(), Error> {
104		let _path = self.tempdir.into_path();
105		self.client.shutdown_immediate().await;
106		self.server.kill()
107			.map_err(ErrorInner::KillServer)?;
108		Ok(())
109	}
110
111	/// Create the temporary directory and spawn a server based on the configuration of the given builder object.
112	async fn from_builder(builder: &TempMongoBuilder) -> Result<Self, Error> {
113		let tempdir = builder.make_temp_dir().map_err(ErrorInner::MakeTempDir)?;
114		let db_dir = tempdir.path().join("db");
115		let socket_path = tempdir.path().join("mongod.sock");
116		let log_path = tempdir.path().join("mongod.log");
117
118		std::fs::create_dir(&db_dir)
119			.map_err(|e| ErrorInner::MakeDbDir(db_dir.clone(), e))?;
120
121		let server = Command::new(builder.get_command())
122			.arg("--bind_ip")
123			.arg(&socket_path)
124			.arg("--dbpath")
125			.arg(db_dir)
126			.arg("--logpath")
127			.arg(&log_path)
128			.arg("--nounixsocket")
129			.arg("--noauth")
130			.spawn()
131			.map_err(|e| ErrorInner::SpawnServer(builder.get_command_string(), e))?;
132		let server = KillOnDrop::new(server);
133
134		let client_options = mongodb::options::ClientOptions::builder()
135			.hosts(vec![mongodb::options::ServerAddress::Unix { path: socket_path.clone() }])
136			.connect_timeout(Duration::from_millis(10))
137			.retry_reads(false)
138			.retry_writes(false)
139			.build();
140		let client = mongodb::Client::with_options(client_options)
141			.map_err(|e| ErrorInner::Connect(socket_path.display().to_string(), e))?;
142
143		client.list_databases(None, None).await
144			.map_err(|e| ErrorInner::Connect(socket_path.display().to_string(), e))?;
145
146		Ok(Self {
147			tempdir,
148			socket_path,
149			log_path,
150			server,
151			client,
152		})
153	}
154}
155
156/// Builder for customizing your [`TempMongo`] object.
157///
158/// After configuring the desirec options, run [`TempMongoBuilder::spawn()`].
159#[derive(Debug)]
160pub struct TempMongoBuilder {
161	/// The parent directory for the temporary directory.
162	///
163	/// Use the system default if set to `None`.
164	parent_directory: Option<PathBuf>,
165
166	/// Clean up the temprorary directory when the [`TempMongo`] object is dropped.
167	clean_on_drop: bool,
168
169	/// The mongdb command to execute.
170	command: Option<OsString>,
171}
172
173impl TempMongoBuilder {
174	/// Create a new builder.
175	pub fn new() -> Self {
176		Self {
177			parent_directory: None,
178			command: None,
179			clean_on_drop: true,
180		}
181	}
182
183	/// Spawn the MongoDB server and connect to it.
184	pub async fn spawn(&self) -> Result<TempMongo, Error> {
185		TempMongo::from_builder(self).await
186	}
187
188	/// Enable or disable cleaning of the temporary state directory when the [`TempMongo`] object is dropped.
189	///
190	/// This can also be changed after creation with [`TempMongo::set_clean_on_drop()`].
191	pub fn clean_on_drop(mut self, clean_on_drop: bool) -> Self {
192		self.clean_on_drop = clean_on_drop;
193		self
194	}
195
196	/// Overwrite the `mongod` command to run.
197	///
198	/// Can be used to run a `mongod` binary from an alternative location.
199	pub fn mongod_command(mut self, command: impl Into<OsString>) -> Self {
200		self.command = Some(command.into());
201		self
202	}
203
204	/// Get the command to execute to run MongoDB.
205	fn get_command(&self) -> &OsStr {
206		self.command
207			.as_deref()
208			.unwrap_or("mongod".as_ref())
209	}
210
211	/// Get the command to execute to run MongDB as a string, for diagnostic purposes.
212	fn get_command_string(&self) -> String {
213		self.get_command().to_string_lossy().into()
214	}
215
216	/// Create a temporary directory according to the configuration of the builder.
217	fn make_temp_dir(&self) -> std::io::Result<TempDir> {
218		match &self.parent_directory {
219			Some(dir) => TempDir::new_in(dir, self.clean_on_drop),
220			None => TempDir::new(self.clean_on_drop),
221		}
222	}
223}
224
225impl Default for TempMongoBuilder {
226	fn default() -> Self {
227		Self::new()
228	}
229}