risedev/task/
mongodb_service.rs

1// Copyright 2026 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::process::Command;
16
17use anyhow::{Context, Result, anyhow};
18
19use super::docker_service::{DockerService, DockerServiceConfig};
20use crate::MongoDbConfig;
21use crate::task::{ExecuteContext, Task};
22
23const REPLICA_SET_NAME: &str = "rs0";
24
25impl DockerServiceConfig for MongoDbConfig {
26    fn id(&self) -> String {
27        self.id.clone()
28    }
29
30    fn is_user_managed(&self) -> bool {
31        self.user_managed
32    }
33
34    fn image(&self) -> String {
35        self.image.clone()
36    }
37
38    fn args(&self) -> Vec<String> {
39        vec![
40            "--replSet".to_owned(),
41            REPLICA_SET_NAME.to_owned(),
42            "--bind_ip_all".to_owned(),
43            "--oplogSize".to_owned(),
44            "128".to_owned(),
45        ]
46    }
47
48    fn ports(&self) -> Vec<(String, String)> {
49        vec![(self.port.to_string(), "27017".to_owned())]
50    }
51
52    fn data_path(&self) -> Option<String> {
53        self.persist_data.then(|| "/data/db".to_owned())
54    }
55}
56
57pub type MongoDbService = DockerService<MongoDbConfig>;
58
59pub struct MongoDbSetupTask {
60    config: MongoDbConfig,
61}
62
63impl MongoDbSetupTask {
64    pub fn new(config: MongoDbConfig) -> Self {
65        Self { config }
66    }
67
68    fn container_name(&self) -> String {
69        format!("risedev-{}", self.config.id)
70    }
71
72    fn mongosh_eval(&self, script: &str) -> Result<std::process::Output> {
73        let output = Command::new("docker")
74            .arg("exec")
75            .arg(self.container_name())
76            .arg("mongosh")
77            .arg("--quiet")
78            .arg("--eval")
79            .arg(script)
80            .output()
81            .context("failed to run mongosh in mongodb container")?;
82        Ok(output)
83    }
84
85    fn wait_replica_set_ready(&self, ctx: &mut ExecuteContext<impl std::io::Write>) -> Result<()> {
86        ctx.pb.set_message("waiting for replica set ready...");
87        ctx.wait(|| {
88            let output = self.mongosh_eval(
89                r#"
90try {
91  const status = rs.status();
92  if (status.ok === 1) {
93    quit(0);
94  }
95  quit(1);
96} catch (err) {
97  if (err.codeName === "NotYetInitialized") {
98    quit(2);
99  }
100  print(err);
101  quit(3);
102}
103"#,
104            )?;
105            match output.status.code() {
106                Some(0) => Ok(()),
107                Some(1) | Some(2) => Err(anyhow!("mongodb replica set is not ready yet")),
108                _ => {
109                    let stderr = String::from_utf8_lossy(&output.stderr);
110                    let stdout = String::from_utf8_lossy(&output.stdout);
111                    Err(anyhow!(
112                        "failed to check mongodb replica set status: {}{}",
113                        stdout.trim(),
114                        stderr.trim()
115                    ))
116                }
117            }
118        })?;
119        Ok(())
120    }
121}
122
123impl Task for MongoDbSetupTask {
124    fn execute(&mut self, ctx: &mut ExecuteContext<impl std::io::Write>) -> Result<()> {
125        if self.config.user_managed {
126            return Ok(());
127        }
128
129        let member = format!("{}:{}", self.config.address, self.config.port);
130        let init_script = format!(
131            r#"
132try {{
133  const status = rs.status();
134  if (status.ok === 1) {{
135    quit(0);
136  }}
137}} catch (err) {{
138  if (err.codeName !== "NotYetInitialized") {{
139    print(err);
140    quit(2);
141  }}
142}}
143
144const result = rs.initiate({{
145  _id: "{REPLICA_SET_NAME}",
146  members: [{{ _id: 0, host: "{member}" }}]
147}});
148
149if (result.ok === 1) {{
150  quit(0);
151}}
152
153printjson(result);
154quit(1);
155"#
156        );
157
158        ctx.pb.set_message("initializing replica set...");
159        ctx.wait(|| {
160            let output = self.mongosh_eval(&init_script)?;
161            if output.status.success() {
162                return Ok(());
163            }
164
165            let stderr = String::from_utf8_lossy(&output.stderr);
166            let stdout = String::from_utf8_lossy(&output.stdout);
167            Err(anyhow!(
168                "failed to initialize mongodb replica set: {}{}",
169                stdout.trim(),
170                stderr.trim()
171            ))
172        })?;
173
174        self.wait_replica_set_ready(ctx)?;
175        Ok(())
176    }
177
178    fn id(&self) -> String {
179        self.config.id.clone()
180    }
181}
risedev/task/mongodb_service.rs

risedev/task/
mongodb_service.rs
