From f211b324c12c00f607845726b2786ca0885ede27 Mon Sep 17 00:00:00 2001
From: Vamshi Maskuri <117595548+varshith257@users.noreply.github.com>
Date: Wed, 11 Dec 2024 02:44:47 +0530
Subject: [PATCH] fix tests

---
 extension/sql/meta.sql               | 38 ++++++++++++++++++----------
 extension/tests/integration_tests.rs | 36 ++++++++++++++++----------
 2 files changed, 46 insertions(+), 28 deletions(-)

diff --git a/extension/sql/meta.sql b/extension/sql/meta.sql
index e66c340..c61743b 100644
--- a/extension/sql/meta.sql
+++ b/extension/sql/meta.sql
@@ -7,30 +7,40 @@ CREATE TABLE vectorize.job (
     last_completion TIMESTAMP WITH TIME ZONE
 );
 
+CREATE INDEX IF NOT EXISTS idx_job_params_table_schema
+ON vectorize.job ((params ->> 'table'), (params ->> 'schema'));
+
 -- create an event trigger function to delete jobs when corresponding tables are dropped
 CREATE OR REPLACE FUNCTION after_drop_trigger()
 RETURNS event_trigger AS $$
 DECLARE
-    dropped_table_name TEXT;
-    dropped_table_schema TEXT;
+    dropped_object RECORD;
+    schema_name TEXT;
+    table_name TEXT;
 BEGIN
-    -- Get the name and schema of the table being dropped
-    FOR dropped_table_name, dropped_table_schema IN
-        SELECT objid::regclass::text, nspname
-        FROM pg_event_trigger_dropped_objects()
-        JOIN pg_class ON objid = pg_class.oid
-        JOIN pg_namespace ON pg_class.relnamespace = pg_namespace.oid
-        WHERE classid = 'pg_class'::regclass
-    LOOP
-        DELETE FROM vectorize.job 
-        WHERE LOWER(params ->> 'table') = LOWER(dropped_table_name)
-          AND LOWER(params ->> 'schema') = LOWER(dropped_table_schema);
+    -- Using pg_event_trigger_dropped_objects() to fetch details of dropped objects.
+    -- This function provides metadata for objects affected by the DROP event.
+    FOR dropped_object IN SELECT * FROM pg_event_trigger_dropped_objects() LOOP
+        IF dropped_object.object_type = 'table' THEN
+            schema_name := split_part(dropped_object.object_identity, '.', 1);
+            table_name := split_part(dropped_object.object_identity, '.', 2);
+            RAISE NOTICE 'Processing drop for table: %, schema: %', table_name, schema_name;
+
+        -- Delete jobs associated with the dropped table
+            DELETE FROM vectorize.job
+            WHERE params ? 'table' AND params ? 'schema'
+              AND params ->> 'table' = table_name
+              AND params ->> 'schema' = schema_name;
+            RAISE NOTICE 'Job deletion executed for table: %, schema: %', table_name, schema_name;
+        END IF;
     END LOOP;
 END;
 $$ LANGUAGE plpgsql;
 
+DROP EVENT TRIGGER IF EXISTS vectorize_job_drop_trigger;
+
 -- create the event trigger for DROP TABLE events
-CREATE EVENT TRIGGER trg_after_drop
+CREATE EVENT TRIGGER vectorize_job_drop_trigger
 ON sql_drop
 WHEN TAG IN ('DROP TABLE')
 EXECUTE FUNCTION after_drop_trigger();
diff --git a/extension/tests/integration_tests.rs b/extension/tests/integration_tests.rs
index 0a60759..a49a259 100644
--- a/extension/tests/integration_tests.rs
+++ b/extension/tests/integration_tests.rs
@@ -58,10 +58,11 @@ async fn test_drop_table_triggers_job_deletion() {
     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!("drop_test_table_{}", test_num);
+    let test_table_name = format!("test_table_{}", test_num);
     let job_name = format!("job_{}", test_num);
 
     common::init_test_table(&test_table_name, &conn).await;
+    common::init_embedding_svc_url(&conn).await;
 
     let create_job_query = format!(
         "SELECT vectorize.table(
@@ -73,40 +74,47 @@ async fn test_drop_table_triggers_job_deletion() {
         );"
     );
     sqlx::query(&create_job_query)
+        .bind(&job_name)
+        .bind(&test_table_name)
         .execute(&conn)
         .await
         .expect("failed to create job");
-
-    // Check row count in vectorize.job before dropping the table
-    let rowcount_before = common::row_count("vectorize.job", &conn).await;
-    assert!(rowcount_before >= 1);
+    
+    // Verify the job exists before dropping the table
+    let job_count_before = common::row_count("vectorize.job", &conn).await;
+    assert_eq!(
+        job_count_before, 1,
+        "Expected 1 job in vectorize.job, found {}",
+        job_count_before
+    );
 
     // Drop the test table with CASCADE to remove dependencies
-    let drop_table_query = format!("DROP TABLE public.{test_table_name} CASCADE;");
-    sqlx::query(&drop_table_query)
+    let drop_query = format!("DROP TABLE public.{test_table_name} CASCADE;");
+    sqlx::query(&drop_query)
         .execute(&conn)
         .await
-        .expect("failed to drop table");
+        .expect("Failed to drop test table");
 
     // Check row count in vectorize.job after dropping the table
-    let rowcount_after = common::row_count("vectorize.job", &conn).await;
+    let job_count_after = common::row_count("vectorize.job", &conn).await;
     assert_eq!(
-        rowcount_after,
-        rowcount_before - 1,
-        "Job was not deleted after table drop"
+        job_count_after, 0,
+        "Job was not deleted after dropping table"
     );
 
     // Verify the specific job no longer exists in vectorize.job
     let job_exists: bool = sqlx::query_scalar(&format!(
-        "SELECT EXISTS (SELECT 1 FROM vectorize.job WHERE name = '{job_name}');"
+        "SELECT EXISTS (SELECT 1 FROM vectorize.job WHERE name = $1);"
     ))
+    .bind(&job_name)
     .fetch_one(&conn)
     .await
     .expect("failed to check job existence");
 
     assert!(
         !job_exists,
-        "Job associated with table `{test_table_name}` was not deleted after table drop"
+        "Job associated with table `{}` was not deleted",
+        test_table_name
     );
 }