From 50f22a6c1701946cf30b241b89750a2c5f00628b Mon Sep 17 00:00:00 2001 From: Erik Mackdanz Date: Fri, 29 Nov 2024 12:57:41 -0600 Subject: [PATCH] DB arg names don't clobber USER --- src/main.rs | 80 +++++++++++++++++++++++++++-------------------------- 1 file changed, 41 insertions(+), 39 deletions(-) diff --git a/src/main.rs b/src/main.rs index 5b0d280..0c2f96d 100644 --- a/src/main.rs +++ b/src/main.rs @@ -19,20 +19,20 @@ //! //! Initialize the database when the database or table doesn't exist already //! -//! Usage: vecsearch init-database [OPTIONS] --password +//! Usage: vecsearch init-database [OPTIONS] --dbpassword //! //! Options: -//! --dbname [default: vsearch] -//! --host [default: localhost] -//! --user [default: cvmigrator] -//! --password [env: PASSWORD=] -//! -h, --help Print help +//! --dbname [default: vsearch] +//! --dbhost [default: localhost] +//! --dbuser [default: cvmigrator] +//! --dbpassword [env: DBPASSWORD=] +//! -h, --help Print help //! ``` //! //! For example: //! //! ```text -//! $ export PASSWORD=$(gpg -d pw-cvmigrator.gpg) +//! $ export DBPASSWORD=$(gpg -d pw-cvmigrator.gpg) //! $ vecsearch init-database //! maybe creating the database //! database vsearch exists already @@ -50,7 +50,7 @@ //! in each invocation. //! //! ```text -//! $ export PASSWORD=$(gpg -d pw-cvmigrator.gpg) +//! $ export DBPASSWORD=$(gpg -d pw-cvmigrator.gpg) //! $ vecsearch index --file testdata/0 --file testdata/1 //! indexing file(s) //! Loaded and encoded 58.565µs @@ -69,7 +69,7 @@ //! superuser). //! //! ```text -//! $ export PASSWORD=$(gpg -d pw-cvmigrator.gpg) +//! $ export DBPASSWORD=$(gpg -d pw-cvmigrator.gpg) //! $ vecsearch search --search feline //! searching for document matches //! Loaded and encoded 49.306µs @@ -84,6 +84,8 @@ //! ## TODO //! //! - why model from PR not main? +//! - init-model command +//! //! use candle_transformers::models::bert::{BertModel, Config, HiddenAct, DTYPE}; @@ -107,13 +109,13 @@ enum Action { dbname: String, #[arg(long,env,default_value="localhost")] - host: String, + dbhost: String, #[arg(long,env,default_value="cvmigrator")] - user: String, + dbuser: String, #[arg(long,env)] - password: String, + dbpassword: String, }, /// Read one document and add it to the database index @@ -123,13 +125,13 @@ enum Action { dbname: String, #[arg(long,env,default_value="localhost")] - host: String, + dbhost: String, #[arg(long,env,default_value="cvmigrator")] - user: String, + dbuser: String, #[arg(long,env)] - password: String, + dbpassword: String, #[arg(long)] /// The file containing document contents. Specify multiple @@ -143,13 +145,13 @@ enum Action { dbname: String, #[arg(long,env,default_value="localhost")] - host: String, + dbhost: String, #[arg(long,env,default_value="cvmigrator")] - user: String, + dbuser: String, #[arg(long,env)] - password: String, + dbpassword: String, #[arg(long)] /// Search for this @@ -227,13 +229,13 @@ impl Args { } } -fn init_database(dbname: String, host: String, user: String, password: String) -> Result<()> { +fn init_database(dbname: String, dbhost: String, dbuser: String, dbpassword: String) -> Result<()> { println!("maybe creating the database"); let mut client = postgres::Config::new() .dbname("postgres") - .host(&host) - .user(&user) - .password(password.clone()) + .host(&dbhost) + .user(&dbuser) + .password(dbpassword.clone()) .connect(NoTls)?; let result = @@ -250,9 +252,9 @@ fn init_database(dbname: String, host: String, user: String, password: String) - println!("maybe creating database objects"); let mut client = postgres::Config::new() .dbname(&dbname) - .host(&host) - .user(&user) - .password(password) + .host(&dbhost) + .user(&dbuser) + .password(dbpassword) .connect(NoTls)?; client.execute("CREATE EXTENSION IF NOT EXISTS vector", &[])?; @@ -290,16 +292,16 @@ fn get_embeddings(input: &String, model: &BertModel, mut tokenizer: Tokenizer) - Ok(embeddings) } -fn index(dbname: String, host: String, user: String, password: String, +fn index(dbname: String, dbhost: String, dbuser: String, dbpassword: String, files: Vec, model: BertModel, tokenizer: Tokenizer) -> Result<()> { println!("indexing file(s)"); let mut client = postgres::Config::new() .dbname(&dbname) - .host(&host) - .user(&user) - .password(password) + .host(&dbhost) + .user(&dbuser) + .password(dbpassword) .connect(NoTls)?; for file in files { @@ -317,7 +319,7 @@ fn index(dbname: String, host: String, user: String, password: String, Ok(()) } -fn search(dbname: String, host: String, user: String, password: String, +fn search(dbname: String, dbhost: String, dbuser: String, dbpassword: String, search: String, model: BertModel, tokenizer: Tokenizer) -> Result<()> { println!("searching for document matches"); @@ -325,9 +327,9 @@ fn search(dbname: String, host: String, user: String, password: String, let mut client = postgres::Config::new() .dbname(&dbname) - .host(&host) - .user(&user) - .password(password) + .host(&dbhost) + .user(&dbuser) + .password(dbpassword) .connect(NoTls)?; for row in client.query("SELECT content FROM documents \ @@ -347,14 +349,14 @@ fn main() -> Result<()> { let (model, tokenizer) = args.build_model_and_tokenizer()?; match args.action { - Action::InitDatabase{ dbname, host, user, password } => { - init_database(dbname, host, user, password)?; + Action::InitDatabase{ dbname, dbhost, dbuser, dbpassword } => { + init_database(dbname, dbhost, dbuser, dbpassword)?; } - Action::Index{ dbname, host, user, password, file } => { - index(dbname, host, user, password, file, model, tokenizer)?; + Action::Index{ dbname, dbhost, dbuser, dbpassword, file } => { + index(dbname, dbhost, dbuser, dbpassword, file, model, tokenizer)?; } - Action::Search{ dbname, host, user, password, search: search_term } => { - search(dbname, host, user, password, search_term, model, tokenizer)?; + Action::Search{ dbname, dbhost, dbuser, dbpassword, search: search_term } => { + search(dbname, dbhost, dbuser, dbpassword, search_term, model, tokenizer)?; } } -- 2.52.0