sqlness_runner/env/
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::collections::HashMap;
16use std::fmt::Display;
17use std::fs::OpenOptions;
18use std::io;
19use std::io::Write;
20use std::path::{Path, PathBuf};
21use std::process::{Child, Command};
22use std::sync::atomic::{AtomicU32, Ordering};
23use std::sync::{Arc, Mutex};
24use std::time::Duration;
25
26use async_trait::async_trait;
27use common_error::ext::ErrorExt;
28use sqlness::{Database, EnvController, QueryContext};
29use tokio::sync::Mutex as TokioMutex;
30
31use crate::client::MultiProtocolClient;
32use crate::cmd::bare::ServerAddr;
33use crate::cmd::compat_case::try_infer_version;
34use crate::formatter::{ErrorFormatter, MysqlFormatter, OutputFormatter, PostgresqlFormatter};
35use crate::protocol_interceptor::{MYSQL, PROTOCOL_KEY};
36use crate::server_mode::{GrpcArgStyle, ServerMode};
37use crate::util;
38use crate::util::{PROGRAM, get_workspace_root, maybe_pull_binary};
39
40// standalone mode
41const SERVER_MODE_STANDALONE_IDX: usize = 0;
42// distributed mode
43const SERVER_MODE_METASRV_IDX: usize = 0;
44const SERVER_MODE_DATANODE_START_IDX: usize = 1;
45const SERVER_MODE_FRONTEND_IDX: usize = 4;
46const SERVER_MODE_FLOWNODE_IDX: usize = 5;
47// Number of datanodes in distributed mode
48const DISTRIBUTED_DATANODE_COUNT: usize = 3;
49
50#[derive(Clone)]
51pub enum WalConfig {
52    RaftEngine,
53    Kafka {
54        /// Indicates whether the runner needs to start a kafka cluster
55        /// (it might be available in the external system environment).
56        needs_kafka_cluster: bool,
57        broker_endpoints: Vec<String>,
58    },
59}
60
61#[derive(Debug, Clone)]
62pub(crate) enum ServiceProvider {
63    Create,
64    External(String),
65}
66
67impl From<&str> for ServiceProvider {
68    fn from(value: &str) -> Self {
69        if value.is_empty() {
70            Self::Create
71        } else {
72            Self::External(value.to_string())
73        }
74    }
75}
76
77#[derive(Clone)]
78pub struct StoreConfig {
79    pub store_addrs: Vec<String>,
80    pub setup_etcd: bool,
81    pub(crate) setup_pg: Option<ServiceProvider>,
82    pub(crate) setup_mysql: Option<ServiceProvider>,
83    pub enable_flat_format: bool,
84    pub enable_gc: bool,
85}
86
87#[derive(Clone)]
88pub struct Env {
89    sqlness_home: PathBuf,
90    server_addrs: ServerAddr,
91    wal: WalConfig,
92
93    /// The path to the directory that contains the pre-built GreptimeDB binary.
94    /// When running in CI, this is expected to be set.
95    /// If not set, this runner will build the GreptimeDB binary itself when needed, and set this field by then.
96    bins_dir: Arc<Mutex<Option<PathBuf>>>,
97    /// The path to the directory that contains the old pre-built GreptimeDB binaries.
98    versioned_bins_dirs: Arc<Mutex<HashMap<String, PathBuf>>>,
99    /// Pull different versions of GreptimeDB on need.
100    pull_version_on_need: bool,
101    /// Store address for metasrv metadata
102    store_config: StoreConfig,
103    /// Extra command line arguments when starting GreptimeDB binaries.
104    extra_args: Vec<String>,
105    /// Cache for the inferred gRPC argument style per `bins_dir`.
106    grpc_arg_style_cache: Arc<Mutex<HashMap<PathBuf, GrpcArgStyle>>>,
107}
108
109#[async_trait]
110impl EnvController for Env {
111    type DB = GreptimeDB;
112
113    async fn start(&self, mode: &str, id: usize, _config: Option<&Path>) -> Self::DB {
114        if self.server_addrs.server_addr.is_some() && id > 0 {
115            panic!("Parallel test mode is not supported when server address is already set.");
116        }
117
118        unsafe {
119            std::env::set_var("SQLNESS_HOME", self.sqlness_home.display().to_string());
120        }
121        match mode {
122            "standalone" => self.start_standalone(id).await,
123            "distributed" => self.start_distributed(id).await,
124            _ => panic!("Unexpected mode: {mode}"),
125        }
126    }
127
128    /// Stop one [`Database`].
129    async fn stop(&self, _mode: &str, mut database: Self::DB) {
130        database.stop();
131    }
132}
133
134impl Env {
135    pub fn new(
136        data_home: PathBuf,
137        server_addrs: ServerAddr,
138        wal: WalConfig,
139        pull_version_on_need: bool,
140        bins_dir: Option<PathBuf>,
141        store_config: StoreConfig,
142        extra_args: Vec<String>,
143    ) -> Self {
144        Self {
145            sqlness_home: data_home,
146            server_addrs,
147            wal,
148            pull_version_on_need,
149            bins_dir: Arc::new(Mutex::new(bins_dir.clone())),
150            versioned_bins_dirs: Arc::new(Mutex::new(HashMap::from_iter([(
151                "latest".to_string(),
152                bins_dir.clone().unwrap_or(util::get_binary_dir("debug")),
153            )]))),
154            store_config,
155            extra_args,
156            grpc_arg_style_cache: Arc::new(Mutex::new(HashMap::new())),
157        }
158    }
159
160    async fn start_standalone(&self, id: usize) -> GreptimeDB {
161        println!("Starting standalone instance id: {id}");
162
163        if self.server_addrs.server_addr.is_some() {
164            self.connect_db(&self.server_addrs, id).await
165        } else {
166            self.build_db();
167            self.setup_wal();
168            let mut db_ctx = GreptimeDBContext::new(self.wal.clone(), self.store_config.clone());
169
170            let server_mode = ServerMode::random_standalone();
171            db_ctx.set_server_mode(server_mode.clone(), SERVER_MODE_STANDALONE_IDX);
172            let server_addr = server_mode.server_addr().unwrap();
173            let server_process = self.start_server(server_mode, &db_ctx, id, true).await;
174
175            let mut greptimedb = self.connect_db(&server_addr, id).await;
176            greptimedb.server_processes = Some(Arc::new(Mutex::new(vec![server_process])));
177            greptimedb.is_standalone = true;
178            greptimedb.ctx = db_ctx;
179
180            greptimedb
181        }
182    }
183
184    async fn start_distributed(&self, id: usize) -> GreptimeDB {
185        self.start_distributed_inner(id).await
186    }
187
188    /// Internal: start a distributed cluster with flownode.
189    async fn start_distributed_inner(&self, id: usize) -> GreptimeDB {
190        if self.server_addrs.server_addr.is_some() {
191            self.connect_db(&self.server_addrs, id).await
192        } else {
193            self.build_db();
194            self.setup_wal();
195            self.setup_etcd();
196            self.setup_pg();
197            self.setup_mysql().await;
198            let mut db_ctx = GreptimeDBContext::new(self.wal.clone(), self.store_config.clone());
199
200            // start a distributed GreptimeDB
201            let meta_server_mode = ServerMode::random_metasrv();
202            let metasrv_port = match &meta_server_mode {
203                ServerMode::Metasrv {
204                    rpc_server_addr, ..
205                } => rpc_server_addr
206                    .split(':')
207                    .nth(1)
208                    .unwrap()
209                    .parse::<u16>()
210                    .unwrap(),
211                _ => panic!(
212                    "metasrv mode not set, maybe running in remote mode which doesn't support restart?"
213                ),
214            };
215            db_ctx.set_server_mode(meta_server_mode.clone(), SERVER_MODE_METASRV_IDX);
216            let meta_server = self.start_server(meta_server_mode, &db_ctx, id, true).await;
217
218            let mut datanodes = Vec::with_capacity(DISTRIBUTED_DATANODE_COUNT);
219            for i in 0..DISTRIBUTED_DATANODE_COUNT {
220                let datanode_mode = ServerMode::random_datanode(metasrv_port, i as u32);
221                db_ctx.set_server_mode(datanode_mode.clone(), SERVER_MODE_DATANODE_START_IDX + i);
222                let datanode = self.start_server(datanode_mode, &db_ctx, id, true).await;
223                datanodes.push(datanode);
224            }
225
226            let frontend_mode = ServerMode::random_frontend(metasrv_port);
227            let server_addr = frontend_mode.server_addr().unwrap();
228            db_ctx.set_server_mode(frontend_mode.clone(), SERVER_MODE_FRONTEND_IDX);
229            let frontend = self.start_server(frontend_mode, &db_ctx, id, true).await;
230
231            let flownode_mode = ServerMode::random_flownode(metasrv_port, 0);
232            db_ctx.set_server_mode(flownode_mode.clone(), SERVER_MODE_FLOWNODE_IDX);
233            let flownode = self.start_server(flownode_mode, &db_ctx, id, true).await;
234
235            let mut greptimedb = self.connect_db(&server_addr, id).await;
236
237            greptimedb.metasrv_process = Some(meta_server).into();
238            greptimedb.server_processes = Some(Arc::new(Mutex::new(datanodes)));
239            greptimedb.frontend_process = Some(frontend).into();
240            greptimedb.flownode_process = Some(flownode).into();
241            greptimedb.is_standalone = false;
242            greptimedb.ctx = db_ctx;
243
244            greptimedb
245        }
246    }
247
248    async fn connect_db(&self, server_addr: &ServerAddr, id: usize) -> GreptimeDB {
249        let grpc_server_addr = server_addr.server_addr.as_ref().unwrap();
250        let pg_server_addr = server_addr.pg_server_addr.as_ref().unwrap();
251        let mysql_server_addr = server_addr.mysql_server_addr.as_ref().unwrap();
252
253        let client =
254            MultiProtocolClient::connect(grpc_server_addr, pg_server_addr, mysql_server_addr).await;
255        GreptimeDB {
256            client: TokioMutex::new(client),
257            server_processes: None,
258            metasrv_process: None.into(),
259            frontend_process: None.into(),
260            flownode_process: None.into(),
261            active_bins_dir: Mutex::new(self.bins_dir.lock().unwrap().clone()),
262            ctx: GreptimeDBContext {
263                time: 0,
264                datanode_id: Default::default(),
265                wal: self.wal.clone(),
266                store_config: self.store_config.clone(),
267                server_modes: Vec::new(),
268            },
269            is_standalone: false,
270            env: self.clone(),
271            id,
272        }
273    }
274
275    fn stop_server(process: &mut Child) {
276        let _ = process.kill();
277        let _ = process.wait();
278    }
279
280    /// Infers which gRPC argument style to use for the binary at `bins_dir`.
281    fn infer_grpc_arg_style(&self, bins_dir: &Path) -> GrpcArgStyle {
282        let cache_key = bins_dir.to_path_buf();
283
284        // Fast path: already cached.
285        {
286            let cache = self.grpc_arg_style_cache.lock().unwrap();
287            if let Some(style) = cache.get(&cache_key) {
288                return *style;
289            }
290        }
291
292        let version = try_infer_version(bins_dir);
293        let style = GrpcArgStyle::for_version(version.as_ref());
294
295        // Insert into cache (may race with another thread, but both detect
296        // the same value, so it's harmless).
297        {
298            let mut cache = self.grpc_arg_style_cache.lock().unwrap();
299            cache.entry(cache_key).or_insert(style);
300        }
301
302        style
303    }
304
305    async fn start_server(
306        &self,
307        mode: ServerMode,
308        db_ctx: &GreptimeDBContext,
309        id: usize,
310        truncate_log: bool,
311    ) -> Child {
312        let bins_dir = self.bins_dir.lock().unwrap().clone().expect(
313            "GreptimeDB binary is not available. Please pass in the path to the directory that contains the pre-built GreptimeDB binary. Or you may call `self.build_db()` beforehand.",
314        );
315
316        self.start_server_with_bins_dir(mode, db_ctx, id, truncate_log, bins_dir)
317            .await
318    }
319
320    async fn start_server_with_bins_dir(
321        &self,
322        mode: ServerMode,
323        db_ctx: &GreptimeDBContext,
324        id: usize,
325        truncate_log: bool,
326        bins_dir: PathBuf,
327    ) -> Child {
328        let log_file_name = match mode {
329            ServerMode::Datanode { node_id, .. } => {
330                db_ctx.incr_datanode_id();
331                format!("greptime-{}-sqlness-datanode-{}.log", id, node_id)
332            }
333            ServerMode::Flownode { .. } => format!("greptime-{}-sqlness-flownode.log", id),
334            ServerMode::Frontend { .. } => format!("greptime-{}-sqlness-frontend.log", id),
335            ServerMode::Metasrv { .. } => format!("greptime-{}-sqlness-metasrv.log", id),
336            ServerMode::Standalone { .. } => format!("greptime-{}-sqlness-standalone.log", id),
337        };
338        let stdout_file_name = self.sqlness_home.join(log_file_name).display().to_string();
339
340        println!("DB instance {id} log file at {stdout_file_name}");
341
342        let stdout_file = OpenOptions::new()
343            .create(true)
344            .write(true)
345            .truncate(truncate_log)
346            .append(!truncate_log)
347            .open(&stdout_file_name)
348            .unwrap();
349
350        let arg_style = self.infer_grpc_arg_style(&bins_dir);
351        let args = mode.get_args(&self.sqlness_home, self, db_ctx, id, arg_style);
352        let check_ip_addrs = mode.check_addrs();
353
354        for check_ip_addr in &check_ip_addrs {
355            if util::check_port(check_ip_addr.parse().unwrap(), Duration::from_secs(1)).await {
356                panic!(
357                    "Port {check_ip_addr} is already in use, please check and retry.",
358                    check_ip_addr = check_ip_addr
359                );
360            }
361        }
362
363        let program = PROGRAM;
364
365        let abs_bins_dir = bins_dir
366            .canonicalize()
367            .expect("Failed to canonicalize bins_dir");
368
369        let mut process = Command::new(abs_bins_dir.join(program))
370            .current_dir(bins_dir.clone())
371            .env("TZ", "UTC")
372            .args(args)
373            .stdout(stdout_file)
374            .spawn()
375            .unwrap_or_else(|error| {
376                panic!(
377                    "Failed to start the DB with subcommand {}, Error: {error}, path: {:?}",
378                    mode.name(),
379                    bins_dir.join(program)
380                );
381            });
382
383        for check_ip_addr in &check_ip_addrs {
384            if !util::check_port(check_ip_addr.parse().unwrap(), Duration::from_secs(30)).await {
385                Env::stop_server(&mut process);
386                panic!(
387                    "{} doesn't up in 30 seconds, check {} for more details.",
388                    mode.name(),
389                    stdout_file_name
390                )
391            }
392        }
393
394        process
395    }
396
397    /// stop and restart the server process
398    pub(crate) async fn restart_server(&self, db: &GreptimeDB, is_full_restart: bool) {
399        let bins_dir = db.active_bins_dir.lock().unwrap().clone().expect(
400            "GreptimeDB binary is not available. Please pass in the path to the directory that contains the pre-built GreptimeDB binary. Or you may call `self.build_db()` beforehand.",
401        );
402
403        {
404            if let Some(server_process) = db.server_processes.clone() {
405                let mut server_processes = server_process.lock().unwrap();
406                for server_process in server_processes.iter_mut() {
407                    Env::stop_server(server_process);
408                }
409            }
410
411            if is_full_restart {
412                if let Some(mut metasrv_process) =
413                    db.metasrv_process.lock().expect("poisoned lock").take()
414                {
415                    Env::stop_server(&mut metasrv_process);
416                }
417                if let Some(mut frontend_process) =
418                    db.frontend_process.lock().expect("poisoned lock").take()
419                {
420                    Env::stop_server(&mut frontend_process);
421                }
422            }
423
424            // Stop flownode if present.
425            if let Some(mut flownode_process) =
426                db.flownode_process.lock().expect("poisoned lock").take()
427            {
428                Env::stop_server(&mut flownode_process);
429            }
430        }
431
432        // check if the server is distributed or standalone
433        let new_server_processes = if db.is_standalone {
434            let server_mode = db
435                .ctx
436                .get_server_mode(SERVER_MODE_STANDALONE_IDX)
437                .cloned()
438                .unwrap();
439            let server_addr = server_mode.server_addr().unwrap();
440            let new_server_process = self
441                .start_server_with_bins_dir(server_mode, &db.ctx, db.id, false, bins_dir.clone())
442                .await;
443
444            let mut client = db.client.lock().await;
445            client
446                .reconnect_mysql_client(&server_addr.mysql_server_addr.unwrap())
447                .await;
448            client
449                .reconnect_pg_client(&server_addr.pg_server_addr.unwrap())
450                .await;
451            vec![new_server_process]
452        } else {
453            db.ctx.reset_datanode_id();
454            if is_full_restart {
455                let metasrv_mode = db
456                    .ctx
457                    .get_server_mode(SERVER_MODE_METASRV_IDX)
458                    .cloned()
459                    .unwrap();
460                let metasrv = self
461                    .start_server_with_bins_dir(
462                        metasrv_mode,
463                        &db.ctx,
464                        db.id,
465                        false,
466                        bins_dir.clone(),
467                    )
468                    .await;
469                db.metasrv_process
470                    .lock()
471                    .expect("lock poisoned")
472                    .replace(metasrv);
473
474                // wait for metasrv to start
475                // since it seems older version of db might take longer to complete election
476                tokio::time::sleep(Duration::from_secs(5)).await;
477            }
478
479            let mut processes = vec![];
480            for i in 0..DISTRIBUTED_DATANODE_COUNT {
481                let datanode_mode = db
482                    .ctx
483                    .get_server_mode(SERVER_MODE_DATANODE_START_IDX + i)
484                    .cloned()
485                    .unwrap();
486                let new_server_process = self
487                    .start_server_with_bins_dir(
488                        datanode_mode,
489                        &db.ctx,
490                        db.id,
491                        false,
492                        bins_dir.clone(),
493                    )
494                    .await;
495                processes.push(new_server_process);
496            }
497
498            if is_full_restart {
499                let frontend_mode = db
500                    .ctx
501                    .get_server_mode(SERVER_MODE_FRONTEND_IDX)
502                    .cloned()
503                    .unwrap();
504                let server_addr = frontend_mode.server_addr().unwrap();
505                let frontend = self
506                    .start_server_with_bins_dir(
507                        frontend_mode,
508                        &db.ctx,
509                        db.id,
510                        false,
511                        bins_dir.clone(),
512                    )
513                    .await;
514                db.frontend_process
515                    .lock()
516                    .expect("lock poisoned")
517                    .replace(frontend);
518
519                // Reconnect protocol clients to the new frontend process
520                // so that MySQL/Postgres queries use the restarted frontend,
521                // not stale connections to the old (killed) process.
522                let mut client = db.client.lock().await;
523                client
524                    .reconnect_mysql_client(server_addr.mysql_server_addr.as_ref().unwrap())
525                    .await;
526                client
527                    .reconnect_pg_client(server_addr.pg_server_addr.as_ref().unwrap())
528                    .await;
529            }
530
531            // Restart flownode.
532            if let Some(flownode_mode) = db.ctx.get_server_mode(SERVER_MODE_FLOWNODE_IDX).cloned() {
533                let flownode = self
534                    .start_server_with_bins_dir(
535                        flownode_mode,
536                        &db.ctx,
537                        db.id,
538                        false,
539                        bins_dir.clone(),
540                    )
541                    .await;
542                db.flownode_process
543                    .lock()
544                    .expect("lock poisoned")
545                    .replace(flownode);
546            }
547
548            processes
549        };
550
551        if let Some(server_processes) = db.server_processes.clone() {
552            let mut server_processes = server_processes.lock().unwrap();
553            *server_processes = new_server_processes;
554        }
555    }
556
557    /// Setup kafka wal cluster if needed. The counterpart is in [GreptimeDB::stop].
558    fn setup_wal(&self) {
559        if matches!(self.wal, WalConfig::Kafka { needs_kafka_cluster, .. } if needs_kafka_cluster) {
560            util::setup_wal();
561        }
562    }
563
564    /// Setup etcd if needed.
565    fn setup_etcd(&self) {
566        if self.store_config.setup_etcd {
567            let client_ports = self
568                .store_config
569                .store_addrs
570                .iter()
571                .map(|s| s.split(':').nth(1).unwrap().parse::<u16>().unwrap())
572                .collect::<Vec<_>>();
573            util::setup_etcd(client_ports, None, None);
574        }
575    }
576
577    /// Setup PostgreSql if needed.
578    fn setup_pg(&self) {
579        if matches!(self.store_config.setup_pg, Some(ServiceProvider::Create)) {
580            let client_ports = self
581                .store_config
582                .store_addrs
583                .iter()
584                .map(|s| s.split(':').nth(1).unwrap().parse::<u16>().unwrap())
585                .collect::<Vec<_>>();
586            let client_port = client_ports.first().unwrap_or(&5432);
587            util::setup_pg(*client_port, None);
588        }
589    }
590
591    /// Setup MySql if needed.
592    async fn setup_mysql(&self) {
593        if matches!(self.store_config.setup_mysql, Some(ServiceProvider::Create)) {
594            let client_ports = self
595                .store_config
596                .store_addrs
597                .iter()
598                .map(|s| s.split(':').nth(1).unwrap().parse::<u16>().unwrap())
599                .collect::<Vec<_>>();
600            let client_port = client_ports.first().unwrap_or(&3306);
601            util::setup_mysql(*client_port, None);
602
603            // Docker of MySQL starts slowly, so we need to wait for a while
604            tokio::time::sleep(Duration::from_secs(10)).await;
605        }
606    }
607
608    /// Build the DB with `cargo build --bin greptime`
609    fn build_db(&self) {
610        let mut bins_dir = self.bins_dir.lock().unwrap();
611        if bins_dir.is_some() {
612            return;
613        }
614
615        println!("Going to build the DB...");
616        let output = Command::new("cargo")
617            .current_dir(util::get_workspace_root())
618            .args([
619                "build",
620                "--bin",
621                "greptime",
622                "--features",
623                "pg_kvbackend,mysql_kvbackend,vector_index",
624            ])
625            .output()
626            .expect("Failed to start GreptimeDB");
627        if !output.status.success() {
628            println!("Failed to build GreptimeDB, {}", output.status);
629            println!("Cargo build stdout:");
630            io::stdout().write_all(&output.stdout).unwrap();
631            println!("Cargo build stderr:");
632            io::stderr().write_all(&output.stderr).unwrap();
633            panic!();
634        }
635
636        bins_dir.replace(util::get_binary_dir("debug"));
637    }
638
639    pub(crate) fn extra_args(&self) -> &Vec<String> {
640        &self.extra_args
641    }
642
643    /// Start a distributed GreptimeDB cluster. Exposed for compat runner.
644    pub(crate) async fn compat_start_distributed(&self, id: usize) -> GreptimeDB {
645        self.start_distributed(id).await
646    }
647
648    /// Full restart of all distributed processes with a new binary directory,
649    /// preserving the same context and data.
650    /// After restart, waits for the frontend gRPC endpoint to become ready.
651    pub(crate) async fn compat_restart_all(&self, db: &GreptimeDB, bins_dir: PathBuf) {
652        *db.active_bins_dir.lock().unwrap() = Some(bins_dir);
653        self.restart_server(db, true).await;
654        self.wait_frontend_ready(db).await;
655    }
656
657    /// Wait for frontend gRPC readiness after restart.
658    async fn wait_frontend_ready(&self, db: &GreptimeDB) {
659        let frontend_mode = db
660            .ctx
661            .get_server_mode(SERVER_MODE_FRONTEND_IDX)
662            .cloned()
663            .unwrap();
664        if let Some(addr) = frontend_mode.check_addrs().first() {
665            println!("Waiting for frontend gRPC readiness at {addr}...");
666            crate::util::retry_with_backoff(
667                || async {
668                    let mut client = db.client.lock().await;
669                    match client.grpc_query("SELECT 1").await {
670                        Ok(_) => Ok(()),
671                        Err(e) => Err(format!("Frontend not ready: {e}")),
672                    }
673                },
674                10,
675                std::time::Duration::from_secs(1),
676            )
677            .await
678            .unwrap_or_else(|e| panic!("Frontend failed to become ready: {e}"));
679        }
680    }
681}
682
683pub struct GreptimeDB {
684    server_processes: Option<Arc<Mutex<Vec<Child>>>>,
685    metasrv_process: Mutex<Option<Child>>,
686    frontend_process: Mutex<Option<Child>>,
687    flownode_process: Mutex<Option<Child>>,
688    client: TokioMutex<MultiProtocolClient>,
689    active_bins_dir: Mutex<Option<PathBuf>>,
690    ctx: GreptimeDBContext,
691    is_standalone: bool,
692    env: Env,
693    id: usize,
694}
695
696impl GreptimeDB {
697    async fn postgres_query(&self, _ctx: QueryContext, query: String) -> Box<dyn Display> {
698        let mut client = self.client.lock().await;
699
700        match client.postgres_query(&query).await {
701            Ok(rows) => Box::new(PostgresqlFormatter::from(rows)),
702            Err(e) => Box::new(e),
703        }
704    }
705
706    async fn mysql_query(&self, _ctx: QueryContext, query: String) -> Box<dyn Display> {
707        let mut client = self.client.lock().await;
708
709        match client.mysql_query(&query).await {
710            Ok(res) => Box::new(MysqlFormatter::from(res)),
711            Err(e) => Box::new(e),
712        }
713    }
714
715    async fn grpc_query(&self, _ctx: QueryContext, query: String) -> Box<dyn Display> {
716        let mut client = self.client.lock().await;
717
718        match client.grpc_query(&query).await {
719            Ok(rows) => Box::new(OutputFormatter::from(rows)),
720            Err(e) => Box::new(ErrorFormatter::from(e)),
721        }
722    }
723
724    /// Handle `QueryContext` directives for compat statement execution.
725    ///
726    /// Inspects `QueryContext` keys set by sqlness interceptors:
727    /// - `restart`: restarts the server (datanode-only) if not using external address.
728    /// - `version`: switches to the specified binary version and performs a full restart.
729    ///
730    /// This does **not** execute queries itself; it only prepares the server state.
731    /// Used by the compat runner.
732    pub(crate) async fn compat_prepare_query_context(&self, ctx: &QueryContext) {
733        if ctx.context.contains_key("restart") && self.env.server_addrs.server_addr.is_none() {
734            self.env.restart_server(self, false).await;
735        } else if let Some(version) = ctx.context.get("version") {
736            let version_bin_dir = self
737                .env
738                .versioned_bins_dirs
739                .lock()
740                .expect("lock poison")
741                .get(version.as_str())
742                .cloned();
743
744            match version_bin_dir {
745                Some(path) if path.join(PROGRAM).is_file() => {
746                    *self.active_bins_dir.lock().unwrap() = Some(path);
747                }
748                _ => {
749                    maybe_pull_binary(version, self.env.pull_version_on_need).await;
750                    let root = get_workspace_root();
751                    let new_path = PathBuf::from_iter([&root, version]);
752                    *self.active_bins_dir.lock().unwrap() = Some(new_path);
753                }
754            }
755
756            self.env.restart_server(self, true).await;
757            // sleep for a while to wait for the server to fully boot up
758            tokio::time::sleep(Duration::from_secs(5)).await;
759        }
760    }
761
762    pub(crate) async fn compat_query(
763        &self,
764        query: &str,
765        ctx: &QueryContext,
766    ) -> Result<String, String> {
767        let mut client = self.client.lock().await;
768
769        // Handle protocol switching
770        if let Some(protocol) = ctx.context.get(PROTOCOL_KEY) {
771            if protocol == MYSQL {
772                return match client.mysql_query(query).await {
773                    Ok(res) => Ok(crate::formatter::MysqlFormatter::from(res).to_string()),
774                    Err(e) => Err(e),
775                };
776            } else {
777                // postgres
778                return match client.postgres_query(query).await {
779                    Ok(rows) => Ok(crate::formatter::PostgresqlFormatter::from(rows).to_string()),
780                    Err(e) => Err(e),
781                };
782            }
783        }
784
785        // Default: gRPC
786        match client.grpc_query(query).await {
787            Ok(output) => Ok(OutputFormatter::from(output).to_string()),
788            Err(e) => {
789                let status_code = e.status_code();
790                let root_cause = e.output_msg();
791                Err(format!(
792                    "Error: {}({status_code}), {root_cause}",
793                    status_code as u32
794                ))
795            }
796        }
797    }
798}
799
800#[async_trait]
801impl Database for GreptimeDB {
802    async fn query(&self, ctx: QueryContext, query: String) -> Box<dyn Display> {
803        if ctx.context.contains_key("restart") && self.env.server_addrs.server_addr.is_none() {
804            self.env.restart_server(self, false).await;
805        } else if let Some(version) = ctx.context.get("version") {
806            let version_bin_dir = self
807                .env
808                .versioned_bins_dirs
809                .lock()
810                .expect("lock poison")
811                .get(version.as_str())
812                .cloned();
813
814            match version_bin_dir {
815                Some(path) if path.join(PROGRAM).is_file() => {
816                    // use version in versioned_bins_dirs
817                    *self.active_bins_dir.lock().unwrap() = Some(path);
818                }
819                _ => {
820                    // use version in dir files
821                    maybe_pull_binary(version, self.env.pull_version_on_need).await;
822                    let root = get_workspace_root();
823                    let new_path = PathBuf::from_iter([&root, version]);
824                    *self.active_bins_dir.lock().unwrap() = Some(new_path);
825                }
826            }
827
828            self.env.restart_server(self, true).await;
829            // sleep for a while to wait for the server to fully boot up
830            tokio::time::sleep(Duration::from_secs(5)).await;
831        }
832
833        if let Some(protocol) = ctx.context.get(PROTOCOL_KEY) {
834            // protocol is bound to be either "mysql" or "postgres"
835            if protocol == MYSQL {
836                self.mysql_query(ctx, query).await
837            } else {
838                self.postgres_query(ctx, query).await
839            }
840        } else {
841            self.grpc_query(ctx, query).await
842        }
843    }
844}
845
846impl GreptimeDB {
847    fn stop(&mut self) {
848        if let Some(server_processes) = self.server_processes.clone() {
849            let mut server_processes = server_processes.lock().unwrap();
850            for mut server_process in server_processes.drain(..) {
851                Env::stop_server(&mut server_process);
852                println!(
853                    "Standalone or Datanode (pid = {}) is stopped",
854                    server_process.id()
855                );
856            }
857        }
858        if let Some(mut metasrv) = self
859            .metasrv_process
860            .lock()
861            .expect("someone else panic when holding lock")
862            .take()
863        {
864            Env::stop_server(&mut metasrv);
865            println!("Metasrv (pid = {}) is stopped", metasrv.id());
866        }
867        if let Some(mut frontend) = self
868            .frontend_process
869            .lock()
870            .expect("someone else panic when holding lock")
871            .take()
872        {
873            Env::stop_server(&mut frontend);
874            println!("Frontend (pid = {}) is stopped", frontend.id());
875        }
876        if let Some(mut flownode) = self
877            .flownode_process
878            .lock()
879            .expect("someone else panic when holding lock")
880            .take()
881        {
882            Env::stop_server(&mut flownode);
883            println!("Flownode (pid = {}) is stopped", flownode.id());
884        }
885        if matches!(self.ctx.wal, WalConfig::Kafka { needs_kafka_cluster, .. } if needs_kafka_cluster)
886        {
887            util::teardown_wal();
888        }
889    }
890
891    /// Stop all processes managed by this GreptimeDB. Exposed for compat runner.
892    pub(crate) fn compat_stop(&mut self) {
893        self.stop();
894    }
895}
896
897impl Drop for GreptimeDB {
898    fn drop(&mut self) {
899        if self.env.server_addrs.server_addr.is_none() {
900            self.stop();
901        }
902    }
903}
904
905pub struct GreptimeDBContext {
906    /// Start time in millisecond
907    time: i64,
908    datanode_id: AtomicU32,
909    wal: WalConfig,
910    store_config: StoreConfig,
911    server_modes: Vec<ServerMode>,
912}
913
914impl GreptimeDBContext {
915    pub fn new(wal: WalConfig, store_config: StoreConfig) -> Self {
916        Self {
917            time: common_time::util::current_time_millis(),
918            datanode_id: AtomicU32::new(0),
919            wal,
920            store_config,
921            server_modes: Vec::new(),
922        }
923    }
924
925    pub(crate) fn time(&self) -> i64 {
926        self.time
927    }
928
929    pub fn is_raft_engine(&self) -> bool {
930        matches!(self.wal, WalConfig::RaftEngine)
931    }
932
933    pub fn kafka_wal_broker_endpoints(&self) -> String {
934        match &self.wal {
935            WalConfig::RaftEngine => String::new(),
936            WalConfig::Kafka {
937                broker_endpoints, ..
938            } => serde_json::to_string(&broker_endpoints).unwrap(),
939        }
940    }
941
942    fn incr_datanode_id(&self) {
943        let _ = self.datanode_id.fetch_add(1, Ordering::Relaxed);
944    }
945
946    fn reset_datanode_id(&self) {
947        self.datanode_id.store(0, Ordering::Relaxed);
948    }
949
950    pub(crate) fn store_config(&self) -> StoreConfig {
951        self.store_config.clone()
952    }
953
954    fn set_server_mode(&mut self, mode: ServerMode, idx: usize) {
955        if idx >= self.server_modes.len() {
956            self.server_modes.resize(idx + 1, mode.clone());
957        }
958        self.server_modes[idx] = mode;
959    }
960
961    fn get_server_mode(&self, idx: usize) -> Option<&ServerMode> {
962        self.server_modes.get(idx)
963    }
964}
sqlness_runner/env/bare.rs

sqlness_runner/env/
bare.rs
