From 466b34b40c20a76d5abc2ff3d28abdd20e1ea162 Mon Sep 17 00:00:00 2001 From: Akhilender Date: Tue, 29 Oct 2024 23:56:41 +0530 Subject: [PATCH 1/6] feat: Add event trigger to handle job cleanup on table drop in vectorize schema Signed-off-by: Akhilender --- extension/sql/meta.sql | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/extension/sql/meta.sql b/extension/sql/meta.sql index 0570668..79d92c7 100644 --- a/extension/sql/meta.sql +++ b/extension/sql/meta.sql @@ -20,6 +20,24 @@ GRANT SELECT ON ALL SEQUENCES IN SCHEMA vectorize TO pg_monitor; ALTER DEFAULT PRIVILEGES IN SCHEMA vectorize GRANT SELECT ON TABLES TO pg_monitor; ALTER DEFAULT PRIVILEGES IN SCHEMA vectorize GRANT SELECT ON SEQUENCES TO pg_monitor; +CREATE OR REPLACE FUNCTION handle_table_drop() +RETURNS event_trigger AS $$ +DECLARE + obj RECORD; +BEGIN + FOR obj IN SELECT * FROM pg_event_trigger_dropped_objects() LOOP + IF obj.object_type = 'table' AND obj.object_schema = 'vectorize' THEN + DELETE FROM vectorize.job + WHERE name = obj.object_name; + END IF; + END LOOP; +END; +$$ LANGUAGE plpgsql; + +CREATE EVENT TRIGGER vectorize_job_drop_trigger +ON sql_drop +WHEN TAG IN ('DROP TABLE') +EXECUTE FUNCTION handle_table_drop(); INSERT INTO vectorize.prompts (prompt_type, sys_prompt, user_prompt) VALUES ( From 25aea008c97de07d3824be1583c9b3b87b2a4c17 Mon Sep 17 00:00:00 2001 From: Akhilender Date: Fri, 15 Nov 2024 11:37:22 +0530 Subject: [PATCH 2/6] fix: made requested changes-1 --- extension/sql/meta.sql | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/extension/sql/meta.sql b/extension/sql/meta.sql index 79d92c7..d98ccf7 100644 --- a/extension/sql/meta.sql +++ b/extension/sql/meta.sql @@ -26,9 +26,10 @@ DECLARE obj RECORD; BEGIN FOR obj IN SELECT * FROM pg_event_trigger_dropped_objects() LOOP - IF obj.object_type = 'table' AND obj.object_schema = 'vectorize' THEN + IF obj.object_type = 'table' THEN DELETE FROM vectorize.job - WHERE name = obj.object_name; + WHERE params ->> 'table' = obj.object_name + AND params ->> 'schema' = obj.object_schema; END IF; END LOOP; END; From 20964f9a4e085d013355383e0298a11bbb059334 Mon Sep 17 00:00:00 2001 From: Akhilender Date: Fri, 15 Nov 2024 12:23:35 +0530 Subject: [PATCH 3/6] fix: added corresponding integration test --- extension/tests/integration_tests.rs | 52 ++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/extension/tests/integration_tests.rs b/extension/tests/integration_tests.rs index 6361682..859bd64 100644 --- a/extension/tests/integration_tests.rs +++ b/extension/tests/integration_tests.rs @@ -860,3 +860,55 @@ async fn test_cohere() { .unwrap(); assert_eq!(search_results.len(), 3); } + +#[ignore] +#[tokio::test] +async fn test_event_trigger_on_table_drop() { + let conn = common::init_database().await; + let mut rng = rand::thread_rng(); + let test_num = rng.gen_range(1..100000); + let test_table_name = format!("products_test_{}", test_num); + let job_name = format!("job_{}", test_num); + + // Initialize the test table and job + common::init_test_table(&test_table_name, &conn).await; + common::init_embedding_svc_url(&conn).await; + + let _ = sqlx::query(&format!( + "SELECT vectorize.table( + job_name => '{job_name}', + \"table\" => '{test_table_name}', + primary_key => 'product_id', + columns => ARRAY['product_name'], + transformer => 'sentence-transformers/all-MiniLM-L6-v2' + );" + )) + .execute(&conn) + .await + .expect("failed to initialize vectorize job"); + + let job_count_before = common::row_count("vectorize.job", &conn).await; + assert_eq!(job_count_before, 1); + + // Drop the test table + let drop_result = sqlx::query(&format!("DROP TABLE {test_table_name};")) + .execute(&conn) + .await; + assert!(drop_result.is_ok(), "Failed to drop the test table"); + + // Verify the job entry is removed from the vectorize.job table + let job_count_after = common::row_count("vectorize.job", &conn).await; + assert_eq!(job_count_after, 0, "Job entry was not removed after table drop"); + + // Attempt to drop a non-associated table and verify no action is taken + let unrelated_table_name = format!("unrelated_test_{}", test_num); + common::init_test_table(&unrelated_table_name, &conn).await; + let _ = sqlx::query(&format!("DROP TABLE {unrelated_table_name};")) + .execute(&conn) + .await + .expect("Failed to drop the unrelated test table"); + + // Ensure vectorize.job is unaffected + let final_job_count = common::row_count("vectorize.job", &conn).await; + assert_eq!(final_job_count, 0, "vectorize.job should remain unaffected by unrelated table drops"); +} \ No newline at end of file From 5ba8c03b964f428b15c108d7fdd917c457c73d4a Mon Sep 17 00:00:00 2001 From: Akhilender Bongirwar <112749383+akhilender-bongirwar@users.noreply.github.com> Date: Mon, 9 Dec 2024 22:52:11 +0530 Subject: [PATCH 4/6] fix: minor-fix in meta.sql file --- extension/sql/meta.sql | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/extension/sql/meta.sql b/extension/sql/meta.sql index d98ccf7..bc608a7 100644 --- a/extension/sql/meta.sql +++ b/extension/sql/meta.sql @@ -20,21 +20,35 @@ GRANT SELECT ON ALL SEQUENCES IN SCHEMA vectorize TO pg_monitor; ALTER DEFAULT PRIVILEGES IN SCHEMA vectorize GRANT SELECT ON TABLES TO pg_monitor; ALTER DEFAULT PRIVILEGES IN SCHEMA vectorize GRANT SELECT ON SEQUENCES TO pg_monitor; -CREATE OR REPLACE FUNCTION handle_table_drop() +CREATE OR REPLACE FUNCTION handle_table_drop() RETURNS event_trigger AS $$ DECLARE obj RECORD; + schema_name TEXT; + table_name TEXT; BEGIN FOR obj IN SELECT * FROM pg_event_trigger_dropped_objects() LOOP IF obj.object_type = 'table' THEN - DELETE FROM vectorize.job - WHERE params ->> 'table' = obj.object_name - AND params ->> 'schema' = obj.object_schema; + schema_name := split_part(obj.object_identity, '.', 1); + table_name := split_part(obj.object_identity, '.', 2); + + RAISE NOTICE 'Event Trigger Fired for table drop: %, schema: %', table_name, schema_name; + + -- Perform cleanup: delete the associated job from the vectorize.job table + DELETE FROM vectorize.job + WHERE params ->> 'table' = table_name + AND params ->> 'schema' = schema_name; + + IF NOT FOUND THEN + RAISE NOTICE 'No matching job found for table: %, schema: %', table_name, schema_name; + END IF; END IF; END LOOP; END; $$ LANGUAGE plpgsql; +DROP EVENT TRIGGER IF EXISTS vectorize_job_drop_trigger; + CREATE EVENT TRIGGER vectorize_job_drop_trigger ON sql_drop WHEN TAG IN ('DROP TABLE') From 91b2306b1d15b22a778faf14ce4303025ca4d7ee Mon Sep 17 00:00:00 2001 From: Akhilender Bongirwar <112749383+akhilender-bongirwar@users.noreply.github.com> Date: Mon, 9 Dec 2024 22:54:24 +0530 Subject: [PATCH 5/6] fix: integration_test test_event_trigger_on_table_drop --- extension/tests/integration_tests.rs | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/extension/tests/integration_tests.rs b/extension/tests/integration_tests.rs index 859bd64..8828b59 100644 --- a/extension/tests/integration_tests.rs +++ b/extension/tests/integration_tests.rs @@ -887,19 +887,30 @@ async fn test_event_trigger_on_table_drop() { .await .expect("failed to initialize vectorize job"); + // Check the job table before dropping the test table let job_count_before = common::row_count("vectorize.job", &conn).await; assert_eq!(job_count_before, 1); // Drop the test table - let drop_result = sqlx::query(&format!("DROP TABLE {test_table_name};")) + let drop_result = sqlx::query(&format!("DROP TABLE {test_table_name} CASCADE;")) .execute(&conn) .await; assert!(drop_result.is_ok(), "Failed to drop the test table"); - // Verify the job entry is removed from the vectorize.job table + // Debug: Check job table after dropping the test table let job_count_after = common::row_count("vectorize.job", &conn).await; assert_eq!(job_count_after, 0, "Job entry was not removed after table drop"); + // Check if the job was deleted + let deleted_job = sqlx::query("SELECT * FROM vectorize.job WHERE params->>'table' = $1 AND params->>'schema' = $2") + .bind(test_table_name) + .bind("public") + .fetch_optional(&conn) + .await + .expect("Failed to fetch job"); + + assert!(deleted_job.is_none(), "Job was not deleted after table drop"); + // Attempt to drop a non-associated table and verify no action is taken let unrelated_table_name = format!("unrelated_test_{}", test_num); common::init_test_table(&unrelated_table_name, &conn).await; @@ -911,4 +922,4 @@ async fn test_event_trigger_on_table_drop() { // Ensure vectorize.job is unaffected let final_job_count = common::row_count("vectorize.job", &conn).await; assert_eq!(final_job_count, 0, "vectorize.job should remain unaffected by unrelated table drops"); -} \ No newline at end of file +} From 4dc3438b4bf063dae6049dd19e8cd2869173a7fa Mon Sep 17 00:00:00 2001 From: Akhilender Bongirwar <112749383+akhilender-bongirwar@users.noreply.github.com> Date: Mon, 9 Dec 2024 23:38:06 +0530 Subject: [PATCH 6/6] fix: remove debugging notices --- extension/sql/meta.sql | 6 ------ 1 file changed, 6 deletions(-) diff --git a/extension/sql/meta.sql b/extension/sql/meta.sql index bc608a7..322a67f 100644 --- a/extension/sql/meta.sql +++ b/extension/sql/meta.sql @@ -32,16 +32,10 @@ BEGIN schema_name := split_part(obj.object_identity, '.', 1); table_name := split_part(obj.object_identity, '.', 2); - RAISE NOTICE 'Event Trigger Fired for table drop: %, schema: %', table_name, schema_name; - -- Perform cleanup: delete the associated job from the vectorize.job table DELETE FROM vectorize.job WHERE params ->> 'table' = table_name AND params ->> 'schema' = schema_name; - - IF NOT FOUND THEN - RAISE NOTICE 'No matching job found for table: %, schema: %', table_name, schema_name; - END IF; END IF; END LOOP; END;