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

sqlness_runner/env/
bare.rs
