Skip to main content

sqlness_runner/cmd/
bare.rs

1// Copyright 2023 Greptime Team
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::path::PathBuf;
16use std::sync::Arc;
17
18use clap::{Parser, ValueEnum};
19use sqlness::interceptor::Registry;
20use sqlness::{ConfigBuilder, Runner};
21
22use crate::cmd::SqlnessConfig;
23use crate::env::bare::{Env, ServiceProvider, StoreConfig, WalConfig};
24use crate::{protocol_interceptor, util};
25
26#[derive(ValueEnum, Debug, Clone)]
27#[clap(rename_all = "snake_case")]
28enum Wal {
29    RaftEngine,
30    Kafka,
31}
32
33// add a group to ensure that all server addresses are set together
34#[derive(clap::Args, Debug, Clone, Default)]
35pub(crate) struct ServerAddr {
36    /// Address of the grpc server.
37    #[clap(short, long)]
38    pub(crate) server_addr: Option<String>,
39
40    /// Address of the postgres server. Must be set if server_addr is set.
41    #[clap(short, long, requires = "server_addr")]
42    pub(crate) pg_server_addr: Option<String>,
43
44    /// Address of the mysql server. Must be set if server_addr is set.
45    #[clap(short, long, requires = "server_addr")]
46    pub(crate) mysql_server_addr: Option<String>,
47}
48
49#[derive(Debug, Parser)]
50/// Run sqlness tests in bare mode.
51pub struct BareCommand {
52    #[clap(flatten)]
53    config: SqlnessConfig,
54
55    /// Addresses of the server.
56    #[command(flatten)]
57    server_addr: ServerAddr,
58
59    /// The type of Wal.
60    #[clap(short, long, default_value = "raft_engine")]
61    wal: Wal,
62
63    /// The kafka wal broker endpoints. This config will suppress sqlness runner
64    /// from starting a kafka cluster, and use the given endpoint as kafka backend.
65    #[clap(short, long)]
66    kafka_wal_broker_endpoints: Option<String>,
67
68    /// The path to the directory where GreptimeDB's binaries resides.
69    /// If not set, sqlness will build GreptimeDB on the fly.
70    #[clap(long)]
71    bins_dir: Option<PathBuf>,
72
73    /// Preserve persistent state in the temporary directory.
74    /// This may affect future test runs.
75    #[clap(long)]
76    preserve_state: bool,
77
78    /// Pull Different versions of GreptimeDB on need.
79    #[clap(long, default_value = "true")]
80    pull_version_on_need: bool,
81
82    /// The store addresses for metadata, if empty, will use memory store.
83    #[clap(long)]
84    store_addrs: Vec<String>,
85
86    /// Whether to setup etcd, by default it is false.
87    #[clap(long, default_value = "false")]
88    setup_etcd: bool,
89
90    /// Whether to setup pg, by default it is false.
91    #[clap(long, default_missing_value = "", num_args(0..=1))]
92    setup_pg: Option<ServiceProvider>,
93
94    /// Whether to setup mysql, by default it is false.
95    #[clap(long, default_missing_value = "", num_args(0..=1))]
96    setup_mysql: Option<ServiceProvider>,
97
98    /// The number of jobs to run in parallel. Default to half of the cores.
99    #[clap(short, long, default_value = "0")]
100    jobs: usize,
101
102    /// Extra command line arguments when starting GreptimeDB binaries.
103    #[clap(long)]
104    extra_args: Vec<String>,
105
106    /// Enable flat format for storage engine (sets default_flat_format = true).
107    #[clap(long, default_value = "false")]
108    enable_flat_format: bool,
109
110    /// Enable garbage collection in metasrv and datanodes.
111    #[clap(long, default_value = "false")]
112    enable_gc: bool,
113}
114
115#[derive(Debug, PartialEq, Eq)]
116struct Parallelism {
117    /// How many environments (standalone, distributed) run at the same time.
118    env: usize,
119    /// How many database instances run at the same time within one environment.
120    instance: usize,
121}
122
123impl BareCommand {
124    fn parallelism(&self) -> Parallelism {
125        let self_managed_kafka =
126            matches!(self.wal, Wal::Kafka) && self.kafka_wal_broker_endpoints.is_none();
127        // External servers are shared by both environments, a filtered run is for
128        // debugging a single case, and each environment starts and tears down the
129        // self-managed Kafka cluster on its own.
130        if self.server_addr.server_addr.is_some()
131            || self.config.test_filter != ".*"
132            || self_managed_kafka
133        {
134            return Parallelism {
135                env: 1,
136                instance: 1,
137            };
138        }
139
140        // Instances of one environment would share the same metadata table and
141        // Kafka topics. Only distributed uses the kv backend, and the two
142        // environments use different Kafka topic prefixes, so the environments
143        // themselves can still run concurrently.
144        if self.setup_etcd
145            || !self.store_addrs.is_empty()
146            || self.setup_pg.is_some()
147            || self.setup_mysql.is_some()
148            || self.kafka_wal_broker_endpoints.is_some()
149        {
150            return Parallelism {
151                env: self.jobs,
152                instance: 1,
153            };
154        }
155
156        Parallelism {
157            env: self.jobs,
158            instance: self.jobs,
159        }
160    }
161
162    pub async fn run(mut self) {
163        let temp_dir = tempfile::Builder::new()
164            .prefix("sqlness")
165            .tempdir()
166            .unwrap();
167        let sqlness_home = temp_dir.keep();
168
169        let mut interceptor_registry: Registry = Default::default();
170        interceptor_registry.register(
171            protocol_interceptor::PREFIX,
172            Arc::new(protocol_interceptor::ProtocolInterceptorFactory),
173        );
174
175        if let Some(d) = &self.config.case_dir
176            && !d.is_dir()
177        {
178            panic!("{} is not a directory", d.display());
179        }
180        if self.jobs == 0 {
181            self.jobs = num_cpus::get() / 2;
182        }
183        let parallelism = self.parallelism();
184        println!(
185            "Environment parallelism: {}, instance parallelism per environment: {}",
186            parallelism.env, parallelism.instance
187        );
188
189        let config = ConfigBuilder::default()
190            .case_dir(util::get_case_dir(self.config.case_dir))
191            .fail_fast(self.config.fail_fast)
192            .test_filter(self.config.test_filter)
193            .follow_links(true)
194            .env_config_file(self.config.env_config_file)
195            .interceptor_registry(interceptor_registry)
196            .parallelism(parallelism.instance)
197            .env_parallelism(parallelism.env)
198            .build()
199            .unwrap();
200
201        let wal = match self.wal {
202            Wal::RaftEngine => WalConfig::RaftEngine,
203            Wal::Kafka => WalConfig::Kafka {
204                needs_kafka_cluster: self.kafka_wal_broker_endpoints.is_none(),
205                broker_endpoints: self
206                    .kafka_wal_broker_endpoints
207                    .map(|s| s.split(',').map(|s| s.to_string()).collect())
208                    // otherwise default to the same port in `kafka-cluster.yml`
209                    .unwrap_or(vec!["127.0.0.1:9092".to_string()]),
210            },
211        };
212
213        let store = StoreConfig {
214            store_addrs: self.store_addrs.clone(),
215            setup_etcd: self.setup_etcd,
216            setup_pg: self.setup_pg,
217            setup_mysql: self.setup_mysql,
218            enable_flat_format: self.enable_flat_format,
219            enable_gc: self.enable_gc,
220        };
221
222        let runner = Runner::new(
223            config,
224            Env::new(
225                sqlness_home.clone(),
226                self.server_addr,
227                wal,
228                self.pull_version_on_need,
229                self.bins_dir,
230                store,
231                self.extra_args,
232            ),
233        );
234        match runner.run().await {
235            Ok(_) => println!("\x1b[32mAll sqlness tests passed!\x1b[0m"),
236            Err(e) => {
237                println!("\x1b[31mTest failed: {}\x1b[0m", e);
238                std::process::exit(1);
239            }
240        }
241
242        // clean up and exit
243        if !self.preserve_state {
244            if self.setup_etcd {
245                println!("Stopping etcd");
246                util::stop_rm_etcd();
247            }
248            // TODO(weny): remove postgre and mysql containers
249            println!("Removing state in {:?}", sqlness_home);
250            tokio::fs::remove_dir_all(sqlness_home).await.unwrap();
251        }
252    }
253}
254
255#[cfg(test)]
256mod tests {
257    use super::*;
258
259    fn parallelism(jobs: &str, args: &[&str]) -> Parallelism {
260        let cmd = BareCommand::try_parse_from(["bare", "-j", jobs].iter().chain(args)).unwrap();
261        cmd.parallelism()
262    }
263
264    #[test]
265    fn test_parallelism() {
266        let serial = Parallelism {
267            env: 1,
268            instance: 1,
269        };
270        let env_only = Parallelism {
271            env: 4,
272            instance: 1,
273        };
274
275        assert_eq!(
276            parallelism("4", &[]),
277            Parallelism {
278                env: 4,
279                instance: 4
280            }
281        );
282        assert_eq!(parallelism("4", &["--setup-pg"]), env_only);
283        assert_eq!(parallelism("4", &["--setup-mysql"]), env_only);
284        assert_eq!(parallelism("4", &["--setup-etcd"]), env_only);
285        assert_eq!(
286            parallelism("4", &["--store-addrs", "127.0.0.1:2379"]),
287            env_only
288        );
289        assert_eq!(
290            parallelism("4", &["-w", "kafka", "-k", "127.0.0.1:9092"]),
291            env_only
292        );
293        assert_eq!(parallelism("4", &["-w", "kafka"]), serial);
294        assert_eq!(parallelism("4", &["-t", "basic"]), serial);
295        assert_eq!(
296            parallelism(
297                "4",
298                &[
299                    "-s",
300                    "127.0.0.1:4001",
301                    "-p",
302                    "127.0.0.1:4003",
303                    "-m",
304                    "127.0.0.1:4002"
305                ]
306            ),
307            serial
308        );
309        assert_eq!(parallelism("1", &["--setup-pg"]), serial);
310    }
311}