address comments: migration path as arg and return values

crates/sui-rpc-benchmark/src/direct/query_generator.rs

@@ -4,8 +4,11 @@
 use sqlparser::ast::{ColumnOption, CreateIndex, CreateTable, Statement, TableConstraint};
 use sqlparser::dialect::PostgreSqlDialect;
 use sqlparser::parser::Parser;
+use std::path::PathBuf;
 
-pub struct QueryGenerator;
+pub struct QueryGenerator {
+    pub migration_path: PathBuf,
+}
 
 #[derive(Debug, Clone)]
 pub struct BenchmarkQuery {
@@ -15,28 +18,22 @@ pub struct BenchmarkQuery {
 }
 
 impl QueryGenerator {
-    fn read_sqls() -> Result<Vec<String>, anyhow::Error> {
-        let current_dir = std::env::current_dir()?;
-        let migration_path = current_dir
-            .parent() // up to crates/
-            .and_then(|p| p.parent()) // up to sui/
-            .map(|p| p.join("crates/sui-indexer-alt/migrations"))
-            .ok_or_else(|| anyhow::anyhow!("Could not find migrations directory"))?;
-        let migration_path = migration_path.to_str().unwrap();
-        let mut sqls = Vec::new();
-
-        Self::read_sql_impl(std::path::Path::new(migration_path), &mut sqls)?;
+    fn read_sqls(&self) -> Result<Vec<String>, anyhow::Error> {
+        let migration_path = self.migration_path.to_str().unwrap();
+        let sqls = Self::read_sql_impl(std::path::Path::new(migration_path))?;
         println!("Read {} up.sql files from migrations directory", sqls.len());
         Ok(sqls)
     }
 
-    fn read_sql_impl(dir: &std::path::Path, sqls: &mut Vec<String>) -> Result<(), anyhow::Error> {
+    fn read_sql_impl(dir: &std::path::Path) -> Result<Vec<String>, anyhow::Error> {
+        let mut sqls: Vec<String> = Vec::new();
         if dir.is_dir() {
             for entry in std::fs::read_dir(dir)? {
                 let entry = entry?;
                 let path = entry.path();
                 if path.is_dir() {
-                    Self::read_sql_impl(&path, sqls)?;
+                    let inner_sqls = Self::read_sql_impl(&path)?;
+                    sqls.extend(inner_sqls);
                 } else if path.is_file()
                     && path
                         .file_name()
@@ -49,11 +46,11 @@ impl QueryGenerator {
                 }
             }
         }
-        Ok(())
+        Ok(sqls)
     }
 
-    pub fn generate_benchmark_queries() -> Result<Vec<BenchmarkQuery>, anyhow::Error> {
-        let sqls = Self::read_sqls()?;
+    pub fn generate_benchmark_queries(&self) -> Result<Vec<BenchmarkQuery>, anyhow::Error> {
+        let sqls = self.read_sqls()?;
         let mut benchmark_queries = Vec::new();
         for sql in sqls {
             let queries = sql_to_benchmark_queries(&sql)?;

crates/sui-rpc-benchmark/src/lib.rs

@@ -3,6 +3,8 @@
 
 pub mod direct;
 
+use std::path::PathBuf;
+
 use anyhow::Result;
 use clap::{Parser, Subcommand};
 
@@ -29,6 +31,11 @@ pub enum Command {
             default_value = "postgres://postgres:postgres@localhost:5432/sui"
         )]
         db_url: String,
+        #[clap(
+            long,
+            default_value = concat!(env!("CARGO_MANIFEST_DIR"), "/../sui-indexer-alt/migrations")
+        )]
+        migration_path: PathBuf,
     },
     /// Benchmark JSON RPC endpoints
     #[clap(name = "jsonrpc")]
@@ -48,9 +55,17 @@ pub async fn run_benchmarks() -> Result<(), anyhow::Error> {
     let opts: Opts = Opts::parse();
 
     match opts.command {
-        Command::DirectQuery { db_url } => {
-            println!("Running direct query benchmark against {}", db_url);
-            let benchmark_queries = QueryGenerator::generate_benchmark_queries()?;
+        Command::DirectQuery {
+            db_url,
+            migration_path,
+        } => {
+            println!(
+                "Running direct query benchmark against DB {} and migrations {}",
+                db_url,
+                migration_path.display()
+            );
+            let query_generator = QueryGenerator { migration_path };
+            let benchmark_queries = query_generator.generate_benchmark_queries()?;
             println!("Generated {} benchmark queries", benchmark_queries.len());
             let query_executor = QueryExecutor::new(db_url.as_str(), benchmark_queries).await?;
             query_executor.run().await?;