Refactor tests (#71)

* fix test init

* fix example create

CONTRIBUTING.md

@@ -175,8 +175,8 @@ The following can be found within the this project's README, under [Hugging Face
 Begin by creating a `producs` table with the dataset that comes included with `pg_vectorize`.
 
 ```sql
-CREATE TABLE products AS
-TABLE vectorize.example_products WITH DATA;
+CREATE TABLE products (LIKE vectorize.example_products INCLUDING ALL);
+INSERT INTO products SELECT * FROM vectorize.example_products;
 ```
 
 You can then confirm everything is correct by running the following:

README.md

@@ -104,8 +104,8 @@ Follow the [installation](#installation) steps if you haven't already.
 Setup a products table. Copy from the example data provided by the extension.
 
 ```sql
-CREATE TABLE products AS 
-TABLE vectorize.example_products WITH DATA;
+CREATE TABLE products (LIKE vectorize.example_products INCLUDING ALL);
+INSERT INTO products SELECT * FROM vectorize.example_products;
 ```
 
 ```sql
@@ -166,8 +166,8 @@ SELECT pg_reload_conf();
 Create an example table if it does not already exist.
 
 ```sql
-CREATE TABLE products AS 
-TABLE vectorize.example_products WITH DATA;
+CREATE TABLE products (LIKE vectorize.example_products INCLUDING ALL);
+INSERT INTO products SELECT * FROM vectorize.example_products;
 ```
 
 Initialize a table for RAG. We'll use an open source Sentence Transformer to generate embeddings.

docs/examples/openai_embeddings.md

@@ -13,8 +13,8 @@ SELECT pg_reload_conf();
 Create an example table if it does not already exist.
 
 ```sql
-CREATE TABLE products AS 
-TABLE vectorize.example_products WITH DATA;
+CREATE TABLE products (LIKE vectorize.example_products INCLUDING ALL);
+INSERT INTO products SELECT * FROM vectorize.example_products;
 ```
 
 Then create the job.

docs/examples/sentence_transformers.md

@@ -34,8 +34,8 @@ SELECT pg_reload_conf();
 Create an example table if it does not already exist.
 
 ```sql
-CREATE TABLE products AS 
-TABLE vectorize.example_products WITH DATA;
+CREATE TABLE products (LIKE vectorize.example_products INCLUDING ALL);
+INSERT INTO products SELECT * FROM vectorize.example_products;
 ```
 
 ```sql

tests/integration_tests.rs

@@ -9,7 +9,8 @@ async fn test_scheduled_job() {
     let conn = common::init_database().await;
     let mut rng = rand::thread_rng();
     let test_num = rng.gen_range(1..100000);
-    let test_table_name = common::init_test_table(test_num, &conn).await;
+    let test_table_name = format!("products_test_{}", test_num);
+    common::init_test_table(&test_table_name, &conn).await;
     let job_name = format!("job_{}", test_num);
 
     common::init_embedding_svc_url(&conn).await;
@@ -59,7 +60,8 @@ async fn test_realtime_job() {
     common::init_embedding_svc_url(&conn).await;
     let mut rng = rand::thread_rng();
     let test_num = rng.gen_range(1..100000);
-    let test_table_name = common::init_test_table(test_num, &conn).await;
+    let test_table_name = format!("products_test_{}", test_num);
+    common::init_test_table(&test_table_name, &conn).await;
     let job_name = format!("job_{}", test_num);
 
     println!("test_table_name: {}", test_table_name);
@@ -124,7 +126,8 @@ async fn test_rag() {
     common::init_embedding_svc_url(&conn).await;
     let mut rng = rand::thread_rng();
     let test_num = rng.gen_range(1..100000);
-    let test_table_name = common::init_test_table(test_num, &conn).await;
+    let test_table_name = format!("products_test_{}", test_num);
+    common::init_test_table(&test_table_name, &conn).await;
     let agent_name = format!("agent_{}", test_num);
 
     println!("test_table_name: {}", test_table_name);
@@ -157,7 +160,8 @@ async fn test_rag_alternate_schema() {
     common::init_embedding_svc_url(&conn).await;
     let mut rng = rand::thread_rng();
     let test_num = rng.gen_range(1..100000);
-    let test_table_name = common::init_test_table(test_num, &conn).await;
+    let test_table_name = format!("products_test_{}", test_num);
+    common::init_test_table(&test_table_name, &conn).await;
     let agent_name = format!("agent_{}", test_num);
     println!("test_table_name: {}", test_table_name);
     println!("agent_name: {}", agent_name);
@@ -191,14 +195,8 @@ async fn test_static() {
     common::init_embedding_svc_url(&conn).await;
     let mut rng = rand::thread_rng();
     let test_table_name = "products_test_static";
+    common::init_test_table(&test_table_name, &conn).await;
     let job_name = "static_test_job";
-    let query = format!(
-        "CREATE TABLE IF NOT EXISTS {test_table_name} AS TABLE vectorize.example_products WITH DATA;"
-    );
-    let _ = sqlx::query(&query)
-        .execute(&conn)
-        .await
-        .expect("failed to create static test table");
 
     // initialize a job
     let _ = sqlx::query(&format!(

tests/util.rs

@@ -84,14 +84,26 @@ pub mod common {
         Ok(options)
     }
 
-    pub async fn init_test_table(test_num: i32, conn: &Pool<Postgres>) -> String {
-        let table = format!("product_{test_num}");
-        let q = format!("CREATE TABLE {table} AS TABLE vectorize.example_products WITH DATA");
-        let _ = sqlx::query(&q)
-            .execute(conn)
-            .await
-            .expect("failed to create test table");
-        table
+    pub async fn init_test_table(table: &str, conn: &Pool<Postgres>) {
+        let create = format!(
+            "CREATE TABLE IF NOT EXISTS {table} (LIKE vectorize.example_products INCLUDING ALL);"
+        );
+        let insert = format!(
+            "
+        DO $$
+        BEGIN
+            IF (SELECT COUNT(*) FROM {table}) = 0 THEN
+                INSERT INTO {table} SELECT * FROM vectorize.example_products;
+            END IF;
+        END $$;
+       "
+        );
+        for q in vec![create, insert] {
+            sqlx::query(&q)
+                .execute(conn)
+                .await
+                .expect("failed to create test table");
+        }
     }
 
     pub async fn row_count(fq_table_name: &str, conn: &Pool<Postgres>) -> i64 {