From 6729d4e6848a384437527a93949b970bc8a27b8f Mon Sep 17 00:00:00 2001
From: wajeht <58354193+wajeht@users.noreply.github.com>
Date: Fri, 16 Aug 2024 10:47:22 -0500
Subject: [PATCH] refactor(utils): Improve runFreshMigration function to drop
 all tables before running fresh migrations

---
 src/handler.ts |  2 +-
 src/utils.ts   | 38 ++++++++++++++++++++++++++++++--------
 2 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/src/handler.ts b/src/handler.ts
index 80f47bd..34e1f77 100644
--- a/src/handler.ts
+++ b/src/handler.ts
@@ -1,6 +1,6 @@
 import { db } from './db/db';
 import { Request, Response } from 'express';
-import { cleanDatabase, runFreshMigration, runMigrations, seedDatabase } from './utils';
+import { cleanDatabase, runFreshMigration, seedDatabase } from './utils';
 
 // GET /healthz
 export function getHealthzHandler(req: Request, res: Response) {
diff --git a/src/utils.ts b/src/utils.ts
index debece2..1a4969a 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -57,20 +57,19 @@ export async function cleanDatabase() {
 	}
 }
 
-export async function runFreshMigration() {
+export async function runFreshMigration(force: boolean = false) {
 	try {
-		if (appConfig.env !== 'production') {
-			console.log('Cannot run fresh migration on non-production environment');
-			return;
-		}
+		// if (appConfig.env === 'production' && !force) {
+		// 	console.log('Cannot run fresh migration on production environment without force flag');
+		// 	return;
+		// }
 
 		const config = {
 			directory: path.resolve(path.join(process.cwd(), 'dist', 'src', 'db', 'migrations')),
 		};
 
-		console.log('Rolling back all migrations...');
-		await db.migrate.rollback(config, true);
-		console.log('All migrations have been rolled back.');
+		console.log('Dropping all tables...');
+		await dropAllTables();
 
 		console.log('Running fresh migrations...');
 		const [batchNo, migrations] = await db.migrate.latest(config);
@@ -93,6 +92,29 @@ export async function runFreshMigration() {
 	}
 }
 
+async function dropAllTables() {
+	const tables = await db.raw(`
+        SELECT tablename FROM pg_tables
+        WHERE schemaname = current_schema()
+        AND tablename != 'knex_migrations'
+        AND tablename != 'knex_migrations_lock'
+    `);
+
+	await db.raw('SET session_replication_role = replica');
+
+	for (const row of tables.rows) {
+		const tableName = row.tablename;
+		try {
+			await db.schema.dropTableIfExists(tableName);
+			console.log(`Dropped table: ${tableName}`);
+		} catch (error) {
+			console.warn(`Failed to drop table ${tableName}:`, error);
+		}
+	}
+
+	await db.raw('SET session_replication_role = DEFAULT');
+}
+
 export async function runMigrations() {
 	try {
 		if (appConfig.env !== 'production') {