Skip to content

Commit

Permalink
fix tests
Browse files Browse the repository at this point in the history
  • Loading branch information
varshith257 authored Dec 10, 2024
1 parent db0a475 commit f211b32
Show file tree
Hide file tree
Showing 2 changed files with 46 additions and 28 deletions.
38 changes: 24 additions & 14 deletions extension/sql/meta.sql
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand Down
36 changes: 22 additions & 14 deletions extension/tests/integration_tests.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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
);
}

Expand Down

0 comments on commit f211b32

Please sign in to comment.