Made collection finding safer (#766)

SilasMarvin · web-flow · commit df9719cb9ee1 · 2023-06-19T19:45:19.000-07:00
diff --git a/pgml-sdks/rust/pgml/src/collection.rs b/pgml-sdks/rust/pgml/src/collection.rs
@@ -751,27 +751,6 @@ impl Collection {
         Ok(results)
     }
 
-    /// Creates a new Collection from a [models::Collection] and a [`PgPool`];
-    pub fn from_model_and_pool(collection_model: models::Collection, pool: PgPool) -> Self {
-        let (
-            documents_table_name,
-            splitters_table_name,
-            models_table_name,
-            transforms_table_name,
-            chunks_table_name,
-        ) = Self::generate_table_names(&collection_model.name);
-        Self {
-            name: collection_model.name,
-            pool,
-            documents_table_name,
-            splitters_table_name,
-            models_table_name,
-            transforms_table_name,
-            chunks_table_name,
-            active: collection_model.active,
-        }
-    }
-
     fn generate_table_names(name: &str) -> (String, String, String, String, String) {
         [
             ".documents",
diff --git a/pgml-sdks/rust/pgml/src/database.rs b/pgml-sdks/rust/pgml/src/database.rs
@@ -80,7 +80,7 @@ impl Database {
         .fetch_optional(self.pool.borrow())
         .await?;
         match collection {
-            Some(c) => Ok(Collection::from_model_and_pool(c, self.pool.clone())),
+            Some(c) => Ok(Collection::new(c.name, self.pool.clone()).await?),
             None => {
                 sqlx::query("INSERT INTO pgml.collections (name) VALUES ($1)")
                     .bind(name)
diff --git a/pgml-sdks/rust/pgml/src/lib.rs b/pgml-sdks/rust/pgml/src/lib.rs
@@ -101,7 +101,7 @@ mod tests {
         let connection_string = env::var("DATABASE_URL").unwrap();
 
         init_logger(LevelFilter::Info).unwrap();
-        let collection_name = "test26";
+        let collection_name = "test27";
 
         let db = Database::new(&connection_string).await.unwrap();
         let collection = db.create_or_get_collection(collection_name).await.unwrap();
