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

sqlness_runner/env/
bare.rs
