From: Erik Mackdanz <erikmack@gmail.com>
Date: Wed, 27 Nov 2024 19:08:10 +0000 (-0600)
Subject: init database implemented
X-Git-Url: https://git.humopery.space/?a=commitdiff_plain;h=3eb998f068b11da1bf8e0db65936ffd3a41f66e8;p=vecsearch.git

init database implemented
---

diff --git a/src/main.rs b/src/main.rs
index 2c25a7f..9bc9364 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -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))", &[])?;
+
 	}
 	_ => {}
     }