59 lines
1.7 KiB
TypeScript
59 lines
1.7 KiB
TypeScript
import { drizzle } from "drizzle-orm/postgres-js";
|
|
import postgres from "postgres";
|
|
import { sql } from "drizzle-orm";
|
|
|
|
const connectionString =
|
|
process.env.DATABASE_URL ??
|
|
"postgresql://sec_cybert:sec_cybert@localhost:5432/sec_cybert";
|
|
|
|
const client = postgres(connectionString);
|
|
const db = drizzle(client);
|
|
|
|
export async function resetTestData() {
|
|
await db.execute(sql`DELETE FROM human_labels WHERE annotator_id LIKE 'e2e-%'`);
|
|
await db.execute(sql`DELETE FROM adjudications WHERE adjudicator_id LIKE 'e2e-%'`);
|
|
await db.execute(sql`DELETE FROM quiz_sessions WHERE annotator_id LIKE 'e2e-%'`);
|
|
await db.execute(sql`DELETE FROM assignments WHERE annotator_id LIKE 'e2e-%'`);
|
|
await db.execute(sql`DELETE FROM annotators WHERE id LIKE 'e2e-%'`);
|
|
}
|
|
|
|
export async function seedE2EData() {
|
|
// Create E2E test annotator
|
|
await db.execute(sql`
|
|
INSERT INTO annotators (id, display_name, password)
|
|
VALUES ('e2e-tester', 'E2E Tester', 'e2e-tester')
|
|
ON CONFLICT DO NOTHING
|
|
`);
|
|
|
|
// Assign 3 paragraphs from the existing data
|
|
await db.execute(sql`
|
|
INSERT INTO assignments (paragraph_id, annotator_id)
|
|
SELECT p.id, 'e2e-tester'
|
|
FROM paragraphs p
|
|
LIMIT 3
|
|
ON CONFLICT DO NOTHING
|
|
`);
|
|
}
|
|
|
|
export async function clearQuizSessions(annotatorId: string) {
|
|
await db.execute(
|
|
sql`DELETE FROM quiz_sessions WHERE annotator_id = ${annotatorId}`,
|
|
);
|
|
}
|
|
|
|
export async function markOnboarded(annotatorId: string) {
|
|
await db.execute(
|
|
sql`UPDATE annotators SET onboarded_at = NOW() WHERE id = ${annotatorId}`,
|
|
);
|
|
}
|
|
|
|
export async function clearOnboarding(annotatorId: string) {
|
|
await db.execute(
|
|
sql`UPDATE annotators SET onboarded_at = NULL WHERE id = ${annotatorId}`,
|
|
);
|
|
}
|
|
|
|
export async function cleanup() {
|
|
await client.end();
|
|
}
|