diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b7c1a31..34a02bb 100644 --- a/CONTRIBUTING.md +++ b/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: diff --git a/README.md b/README.md index e2b611f..308c98e 100644 --- a/README.md +++ b/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. diff --git a/docs/examples/openai_embeddings.md b/docs/examples/openai_embeddings.md index 7924e66..f3afdda 100644 --- a/docs/examples/openai_embeddings.md +++ b/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. diff --git a/docs/examples/sentence_transformers.md b/docs/examples/sentence_transformers.md index bcf95a2..053c89c 100644 --- a/docs/examples/sentence_transformers.md +++ b/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 diff --git a/tests/integration_tests.rs b/tests/integration_tests.rs index 03e0236..19c0e02 100644 --- a/tests/integration_tests.rs +++ b/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!( diff --git a/tests/util.rs b/tests/util.rs index 2b0f243..743124f 100644 --- a/tests/util.rs +++ b/tests/util.rs @@ -84,14 +84,26 @@ pub mod common { Ok(options) } - pub async fn init_test_table(test_num: i32, conn: &Pool) -> 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) { + 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) -> i64 {