]> Humopery - vecsearch.git/commitdiff
init database implemented
authorErik Mackdanz <erikmack@gmail.com>
Wed, 27 Nov 2024 19:08:10 +0000 (13:08 -0600)
committerErik Mackdanz <erikmack@gmail.com>
Wed, 27 Nov 2024 19:08:10 +0000 (13:08 -0600)
src/main.rs

index 2c25a7f0eb313f4f65d9b4d6611c2d387dd070b9..9bc93645e2e3bcbe415ba3ad497ec6ada3daef1f 100644 (file)
@@ -14,6 +14,9 @@ enum Action {
     /// Initialize the database when the table doesn't exist already
     InitDatabase {
 
+       #[arg(long)]
+       dbname: String,
+
        #[arg(long)]
        host: String,
 
@@ -140,14 +143,40 @@ fn main() -> Result<()> {
     }
 
     match args.action {
-       Action::InitDatabase{ host, user, password } => {
-           println!("initializing the database with user {}", user);
-           let _client = postgres::Config::new()
+       Action::InitDatabase{ dbname, host, user, password } => {
+           println!("maybe creating the database");
+           let mut client = postgres::Config::new()
                .dbname("postgres")
                .host(&host.to_string())
                .user(&user.to_string())
+               .password(password.clone())
+               .connect(NoTls)?;
+
+           let result =
+               client.query("SELECT 1 FROM pg_database WHERE datname=$1",&[&dbname])?;
+
+           if result.is_empty() {
+               println!("creating database {}", dbname);
+               client.execute(&format!("create database {}",dbname),&[])?;
+           } else {
+               println!("database {} exists already", dbname);
+           }
+           let _ = client.close();
+
+           println!("maybe creating database objects");
+           let mut client = postgres::Config::new()
+               .dbname(&dbname)
+               .host(&host.to_string())
+               .user(&user.to_string())
                .password(password)
                .connect(NoTls)?;
+
+           client.execute("CREATE EXTENSION IF NOT EXISTS vector", &[])?;
+           client.execute(
+               "CREATE TABLE if not exists documents \
+                (id serial PRIMARY KEY, content text, \
+                embedding vector(384))", &[])?;
+
        }
        _ => {}
     }