Skip to content

Commit 70334c6

Browse files
SilasMarvinSilasMarvin
authored
Added method to check if collection exists (#809)
1 parent f35a920 commit 70334c6

File tree

7 files changed

+59
-37
lines changed

7 files changed

+59
-37
lines changed

‎pgml-sdks/rust/pgml-macros/src/javascript.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -360,6 +360,7 @@ fn get_neon_type(ty: &SupportedType) -> syn::Type {
360360
SupportedType::str | SupportedType::String => {
361361
syn::parse_str("neon::types::JsString").unwrap()
362362
}
363+
SupportedType::bool => syn::parse_str("neon::types::JsBoolean").unwrap(),
363364
SupportedType::Vec(_v) => syn::parse_str("neon::types::JsArray").unwrap(),
364365
SupportedType::S => syn::parse_str("neon::types::JsObject").unwrap(),
365366
SupportedType::Tuple(_t) => syn::parse_str("neon::types::JsObject").unwrap(),
@@ -410,6 +411,7 @@ fn get_typescript_type(ty: &SupportedType) -> String {
410411
match ty {
411412
SupportedType::Reference(r) => get_typescript_type(r),
412413
SupportedType::str | SupportedType::String => "string".to_string(),
414+
SupportedType::bool => "boolean".to_string(),
413415
SupportedType::Option(o) => get_typescript_type(o),
414416
SupportedType::Vec(v) => format!("{}[]", get_typescript_type(v)),
415417
SupportedType::HashMap((k, v)) => {

‎pgml-sdks/rust/pgml-macros/src/python.rs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -366,6 +366,7 @@ fn get_python_type(ty: &SupportedType) -> String {
366366
SupportedType::Reference(r) => get_python_type(r),
367367
SupportedType::S => "Self".to_string(),
368368
SupportedType::str | SupportedType::String => "str".to_string(),
369+
SupportedType::bool => "bool".to_string(),
369370
SupportedType::Option(o) => format!(
370371
"Optional[{}] = {}",
371372
get_python_type(o),

‎pgml-sdks/rust/pgml-macros/src/types.rs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ pub enum SupportedType {
99
Reference(Box<SupportedType>),
1010
str,
1111
String,
12+
bool,
1213
Vec(Box<SupportedType>),
1314
HashMap((Box<SupportedType>, Box<SupportedType>)),
1415
Option(Box<SupportedType>),
@@ -34,8 +35,9 @@ impl ToString for SupportedType {
3435
match self {
3536
SupportedType::Reference(t) => format!("&{}", t.to_string()),
3637
SupportedType::str => "str".to_string(),
37-
SupportedType::Json => "Json".to_string(),
3838
SupportedType::String => "String".to_string(),
39+
SupportedType::bool => "bool".to_string(),
40+
SupportedType::Json => "Json".to_string(),
3941
SupportedType::Vec(v) => format!("Vec<{}>", v.to_string()),
4042
SupportedType::HashMap((k, v)) => {
4143
format!("HashMap<{},{}>", k.to_string(), v.to_string())
@@ -138,6 +140,7 @@ impl<'ast> Visit<'ast> for GetSupportedType {
138140
self.ty = match segment_name.as_str() {
139141
"str" => Some(SupportedType::str),
140142
"String" => Some(SupportedType::String),
143+
"bool" => Some(SupportedType::bool),
141144
"Vec" => Some(SupportedType::Vec(Box::new(
142145
Self::get_type_from_path_argument(&i.arguments),
143146
))),

‎pgml-sdks/rust/pgml/javascript/package-lock.json

Lines changed: 0 additions & 32 deletions
This file was deleted.

‎pgml-sdks/rust/pgml/python/tests/test.py

Lines changed: 9 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,18 +8,25 @@ async def main():
88
collection_name = "ptest22"
99
db = pgml.Database(CONNECTION_STRING)
1010
collection = await db.create_or_get_collection(collection_name)
11+
print("The Collection")
1112
print(collection)
13+
collection_does_exist = await db.does_collection_exist(collection_name)
14+
print("Collection does exist")
15+
print(collection_does_exist)
1216
x = [{'id': '5733be284776f41900661182', 'text': 'Architecturally, the school has a Catholic character. Atop the Main Building\'s gold dome is a golden statue of the Virgin Mary. Immediately in front of the Main Building and facing it, is a copper statue of Christ with arms upraised with the legend "Venite Ad Me Omnes". Next to the Main Building is the Basilica of the Sacred Heart. Immediately behind the basilica is the Grotto, a Marian place of prayer and reflection. It is a replica of the grotto at Lourdes, France where the Virgin Mary reputedly appeared to Saint Bernadette Soubirous in 1858. At the end of the main drive (and in a direct line that connects through 3 statues and the Gold Dome), is a simple, modern stone statue of Mary.', 'title': 'University_of_Notre_Dame'}]
1317
await collection.upsert_documents(x)
1418
await collection.register_text_splitter("recursive_character", {"chunk_size": 1500, "chunk_overlap": 40})
1519
splitters = await collection.get_text_splitters()
20+
print("The Splitters")
1621
print(splitters)
1722
await collection.generate_chunks()
1823
await collection.register_model("embedding", "intfloat/e5-small")
1924
models = await collection.get_models()
25+
print("The Models")
2026
print(models)
2127
await collection.generate_embeddings()
2228
results = await collection.vector_search("small")
29+
print("The Results")
2330
print(results)
2431
await db.archive_collection(collection_name)
2532

@@ -51,5 +58,5 @@ async def query_builder():
5158

5259

5360
if __name__ == "__main__":
54-
asyncio.run(query_builder())
55-
# asyncio.run(main())
61+
# asyncio.run(query_builder())
62+
asyncio.run(main())

‎pgml-sdks/rust/pgml/src/database.rs

Lines changed: 30 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ pub struct Database {
1919
pub pool: PgPool,
2020
}
2121

22-
#[custom_methods(new, create_or_get_collection, archive_collection)]
22+
#[custom_methods(new, create_or_get_collection, does_collection_exist, archive_collection)]
2323
impl Database {
2424
/// Create a new [Database]
2525
///
@@ -85,6 +85,35 @@ impl Database {
8585
}
8686
}
8787

88+
/// Check if a [Collection] exists
89+
///
90+
/// # Arguments
91+
///
92+
/// * `name` - The name of the [Collection]
93+
///
94+
/// # Example
95+
/// ```
96+
/// use pgml::Database;
97+
///
98+
/// const CONNECTION_STRING: &str = "postgres://postgres@localhost:5432/pgml_development";
99+
///
100+
/// async fn example() -> anyhow::Result<()> {
101+
/// let db = Database::new(CONNECTION_STRING).await?;
102+
/// let collection_exists = db.does_collection_exist("collection number 1").await?;
103+
/// // Do stuff with your new found information
104+
/// Ok(())
105+
/// }
106+
/// ```
107+
pub async fn does_collection_exist(&self, name: &str) -> anyhow::Result<bool> {
108+
let collection: Option<models::Collection> = sqlx::query_as::<_, models::Collection>(
109+
"SELECT * FROM pgml.collections WHERE name = $1 AND active = TRUE;",
110+
)
111+
.bind(name)
112+
.fetch_optional(self.pool.borrow())
113+
.await?;
114+
Ok(collection.is_some())
115+
}
116+
88117
/// Archive a [Collection]
89118
///
90119
/// # Arguments

‎pgml-sdks/rust/pgml/src/lib.rs

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,18 @@ mod tests {
9999
Database::new(&connection_string).await.unwrap();
100100
}
101101

102+
103+
#[tokio::test]
104+
async fn can_create_collection() {
105+
let connection_string = env::var("DATABASE_URL").unwrap();
106+
let collection_name = "rctest0";
107+
let db = Database::new(&connection_string).await.unwrap();
108+
let _ = db.create_or_get_collection(collection_name).await.unwrap();
109+
let does_collection_exist = db.does_collection_exist(collection_name).await.unwrap();
110+
assert_eq!(does_collection_exist, true);
111+
db.archive_collection(collection_name).await.unwrap();
112+
}
113+
102114
#[tokio::test]
103115
async fn can_create_collection_and_vector_search() {
104116
let connection_string = env::var("DATABASE_URL").unwrap();
@@ -157,7 +169,7 @@ mod tests {
157169
let collection = db.create_or_get_collection(collection_name).await.unwrap();
158170

159171
let mut documents: Vec<Json> = Vec::new();
160-
for i in 0..200 {
172+
for i in 0..2 {
161173
documents.push(serde_json::json!({
162174
"id": i,
163175
"text": format!("{} This is some document with some filler text filler filler filler filler filler filler filler filler filler", i)

0 commit comments

Comments
 (0)