feat: refactoring complete?

src/storage/db.rs

@@ -1,8 +1,67 @@
-use surrealdb::{engine::remote::ws::Client, Surreal};
-
 use crate::error::ProcessingError;
 
 use super::types::StoredObject;
+use std::ops::Deref;
+use surrealdb::{
+    engine::remote::ws::{Client, Ws},
+    opt::auth::Root,
+    Error, Surreal,
+};
+
+#[derive(Clone)]
+pub struct SurrealDbClient {
+    pub client: Surreal<Client>,
+}
+
+impl SurrealDbClient {
+    /// # Initialize a new datbase client
+    ///
+    /// # Arguments
+    ///
+    /// # Returns
+    /// * `SurrealDbClient` initialized
+    pub async fn new() -> Result<Self, Error> {
+        let db = Surreal::new::<Ws>("127.0.0.1:8000").await?;
+
+        // Sign in to database
+        db.signin(Root {
+            username: "root_user",
+            password: "root_password",
+        })
+        .await?;
+
+        // Set namespace
+        db.use_ns("test").use_db("test").await?;
+
+        Ok(SurrealDbClient { client: db })
+    }
+
+    pub async fn rebuild_indexes(&self) -> Result<(), Error> {
+        self.client
+            .query("REBUILD INDEX IF EXISTS idx_embedding ON text_chunk")
+            .await?;
+        self.client
+            .query("REBUILD INDEX IF EXISTS embeddings ON knowledge_entity")
+            .await?;
+        Ok(())
+    }
+
+    pub async fn drop_table<T>(&self) -> Result<(), Error>
+    where
+        T: StoredObject + Send + Sync + 'static,
+    {
+        let _deleted: Vec<T> = self.client.delete(T::table_name()).await?;
+        Ok(())
+    }
+}
+
+impl Deref for SurrealDbClient {
+    type Target = Surreal<Client>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.client
+    }
+}
 
 /// Operation to store a object in SurrealDB, requires the struct to implement StoredObject
 ///