risingwave_regress_test/
psql.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use anyhow::{Context, bail};
16use tokio::process::Command;
17use tracing::{debug, info};
18
19use crate::Opts;
20
21const PG_DB_NAME: &str = "postgres";
22
23pub(crate) struct Psql {
24    opts: Opts,
25}
26
27pub(crate) struct PsqlCommandBuilder {
28    database: String,
29    cmd: Command,
30}
31
32impl Psql {
33    pub(crate) fn new(opts: Opts) -> Self {
34        Self { opts }
35    }
36
37    #[allow(clippy::unused_async)]
38    pub(crate) async fn init(&self) -> anyhow::Result<()> {
39        info!("Initializing instances.");
40
41        for _db in [self.opts.database_name(), PG_DB_NAME] {
42            // self.drop_database_if_exists(db).await?;
43            // self.create_database(db).await?;
44        }
45
46        Ok(())
47    }
48
49    pub(crate) async fn create_database<S: AsRef<str>>(&self, db: S) -> anyhow::Result<()> {
50        info!("Creating database {}", db.as_ref());
51
52        let mut cmd = PsqlCommandBuilder::new(PG_DB_NAME, &self.opts)
53            .add_cmd(format!(
54                r#"CREATE DATABASE "{}" TEMPLATE=template0 LC_COLLATE='C' LC_CTYPE='C'"#,
55                db.as_ref()
56            ))
57            .build();
58
59        let status = cmd
60            .status()
61            .await
62            .with_context(|| format!("Failed to execute command: {:?}", cmd))?;
63        if status.success() {
64            info!("Succeeded to create database {}", db.as_ref());
65            Ok(())
66        } else {
67            bail!("Failed to create database {}", db.as_ref())
68        }
69    }
70
71    pub(crate) async fn drop_database_if_exists<S: AsRef<str>>(&self, db: S) -> anyhow::Result<()> {
72        info!("Dropping database {} if exists", db.as_ref());
73
74        let mut cmd = PsqlCommandBuilder::new("postgres", &self.opts)
75            .add_cmd(format!(r#"DROP DATABASE IF EXISTS "{}""#, db.as_ref()))
76            .build();
77
78        debug!("Dropping database command is: {:?}", cmd);
79
80        let status = cmd
81            .status()
82            .await
83            .with_context(|| format!("Failed to execute command: {:?}", cmd))?;
84
85        if status.success() {
86            info!("Succeeded to drop database {}", db.as_ref());
87            Ok(())
88        } else {
89            bail!("Failed to drop database {}", db.as_ref())
90        }
91    }
92}
93
94impl PsqlCommandBuilder {
95    pub(crate) fn new<S: ToString>(database: S, opts: &Opts) -> Self {
96        let mut cmd = Command::new("psql");
97        cmd.arg("-X")
98            .args(["-h", opts.host().as_str()])
99            .args(["-p", format!("{}", opts.port()).as_str()]);
100
101        Self {
102            database: database.to_string(),
103            cmd,
104        }
105    }
106
107    pub(crate) fn add_cmd<S: AsRef<str>>(mut self, cmd: S) -> Self {
108        let cmd = cmd.as_ref();
109        let mut escaped_cmd = "".to_owned();
110
111        // Escape any shell double-quote metacharacters
112        for c in cmd.chars() {
113            if r#"\"$`"#.contains(c) {
114                escaped_cmd.push('\\');
115            }
116            escaped_cmd.push(c);
117        }
118
119        // Append command
120        self.cmd.args(["-c", &escaped_cmd]);
121
122        self
123    }
124
125    pub(crate) fn build(mut self) -> Command {
126        self.cmd.arg(&self.database);
127        self.cmd
128    }
129}
risingwave_regress_test/psql.rs

risingwave_regress_test/
psql.rs
