risingwave_regress_test/
psql.rs

1// Copyright 2022 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    pub(crate) fn init(&self) -> anyhow::Result<()> {
38        info!("Initializing instances.");
39
40        for _db in [self.opts.database_name(), PG_DB_NAME] {
41            // self.drop_database_if_exists(db).await?;
42            // self.create_database(db).await?;
43        }
44
45        Ok(())
46    }
47
48    pub(crate) async fn create_database<S: AsRef<str>>(&self, db: S) -> anyhow::Result<()> {
49        info!("Creating database {}", db.as_ref());
50
51        let mut cmd = PsqlCommandBuilder::new(PG_DB_NAME, &self.opts)
52            .add_cmd(format!(
53                r#"CREATE DATABASE "{}" TEMPLATE=template0 LC_COLLATE='C' LC_CTYPE='C'"#,
54                db.as_ref()
55            ))
56            .build();
57
58        let status = cmd
59            .status()
60            .await
61            .with_context(|| format!("Failed to execute command: {:?}", cmd))?;
62        if status.success() {
63            info!("Succeeded to create database {}", db.as_ref());
64            Ok(())
65        } else {
66            bail!("Failed to create database {}", db.as_ref())
67        }
68    }
69
70    pub(crate) async fn drop_database_if_exists<S: AsRef<str>>(&self, db: S) -> anyhow::Result<()> {
71        info!("Dropping database {} if exists", db.as_ref());
72
73        let mut cmd = PsqlCommandBuilder::new("postgres", &self.opts)
74            .add_cmd(format!(r#"DROP DATABASE IF EXISTS "{}""#, db.as_ref()))
75            .build();
76
77        debug!("Dropping database command is: {:?}", cmd);
78
79        let status = cmd
80            .status()
81            .await
82            .with_context(|| format!("Failed to execute command: {:?}", cmd))?;
83
84        if status.success() {
85            info!("Succeeded to drop database {}", db.as_ref());
86            Ok(())
87        } else {
88            bail!("Failed to drop database {}", db.as_ref())
89        }
90    }
91}
92
93impl PsqlCommandBuilder {
94    pub(crate) fn new<S: ToString>(database: S, opts: &Opts) -> Self {
95        let mut cmd = Command::new("psql");
96        cmd.arg("-X")
97            .args(["-h", opts.host().as_str()])
98            .args(["-p", format!("{}", opts.port()).as_str()]);
99
100        Self {
101            database: database.to_string(),
102            cmd,
103        }
104    }
105
106    pub(crate) fn add_cmd<S: AsRef<str>>(mut self, cmd: S) -> Self {
107        let cmd = cmd.as_ref();
108        let mut escaped_cmd = "".to_owned();
109
110        // Escape any shell double-quote metacharacters
111        for c in cmd.chars() {
112            if r#"\"$`"#.contains(c) {
113                escaped_cmd.push('\\');
114            }
115            escaped_cmd.push(c);
116        }
117
118        // Append command
119        self.cmd.args(["-c", &escaped_cmd]);
120
121        self
122    }
123
124    pub(crate) fn build(mut self) -> Command {
125        self.cmd.arg(&self.database);
126        self.cmd
127    }
128}
risingwave_regress_test/psql.rs

risingwave_regress_test/
psql.rs
