sqlsmith_reducer/
reducer.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 std::process::Command;
16use std::time::Duration;
17
18use clap::Parser;
19use risingwave_sqlsmith::reducer::shrink_file;
20use thiserror_ext::AsReport;
21use tokio_postgres::NoTls;
22use tracing_subscriber::EnvFilter;
23
24/// Reduce an sql query
25#[derive(Parser, Debug)]
26#[command(author, version, about, long_about = None)]
27struct Args {
28    /// Input file
29    #[arg(short, long)]
30    input_file: String,
31
32    /// Output file
33    #[arg(short, long)]
34    output_file: String,
35
36    /// Command to restore RW
37    #[clap(long)]
38    run_rw_cmd: String,
39}
40
41#[tokio::main(flavor = "multi_thread", worker_threads = 5)]
42async fn main() {
43    _ = tracing_subscriber::fmt()
44        .with_env_filter(EnvFilter::from_default_env())
45        .with_ansi(console::colors_enabled_stderr() && console::colors_enabled())
46        .with_writer(std::io::stderr)
47        .try_init();
48
49    let args = Args::parse();
50
51    // Execute restore command before connecting to database
52    tracing::info!("Executing restore command: {}", args.run_rw_cmd);
53    let status = Command::new("sh").arg("-c").arg(&args.run_rw_cmd).status();
54
55    match status {
56        Ok(s) if s.success() => tracing::info!("Restore command executed successfully"),
57        Ok(s) => {
58            tracing::error!("Restore command failed with status: {}", s);
59            panic!("Failed to restore RW");
60        }
61        Err(err) => {
62            tracing::error!("Failed to execute restore command: {}", err.as_report());
63            panic!("Failed to execute restore command: {}", err.as_report());
64        }
65    }
66
67    let (client, connection) = tokio_postgres::Config::new()
68        .host("localhost")
69        .port(4566)
70        .dbname("dev")
71        .user("root")
72        .password("")
73        .connect_timeout(Duration::from_secs(5))
74        .connect(NoTls)
75        .await
76        .unwrap_or_else(|e| panic!("Failed to connect to database: {}", e.as_report()));
77
78    tokio::spawn(async move {
79        if let Err(e) = connection.await {
80            tracing::error!(error = %e.as_report(), "Postgres connection error");
81        }
82    });
83
84    shrink_file(
85        &args.input_file,
86        &args.output_file,
87        client,
88        &args.run_rw_cmd,
89    )
90    .await
91    .unwrap();
92}
sqlsmith_reducer/reducer.rs

sqlsmith_reducer/
reducer.rs
