From 5c18c76f8bbb8c4ac1c31367390ef703f044f634 Mon Sep 17 00:00:00 2001
From: Julien Ripouteau <julien@ripouteau.com>
Date: Fri, 23 Feb 2024 13:01:50 +0100
Subject: [PATCH] refactor: change sqlite dropAllTables implementation (#1001)

* refactor: change dropAllTables implementation

* test: add test for dropTables with foreign keys constraint

* refactor: handle pragma foreign_keys

* fix: test for mssql

* test: fix for mssql
---
 src/dialects/base_sqlite.ts       | 36 ++++++++++++++------
 test-helpers/index.ts             |  6 ----
 test/database/drop_tables.spec.ts | 55 +++++++++++++++++++++++++++++++
 3 files changed, 81 insertions(+), 16 deletions(-)

diff --git a/src/dialects/base_sqlite.ts b/src/dialects/base_sqlite.ts
index d81271e4..23aa1f94 100644
--- a/src/dialects/base_sqlite.ts
+++ b/src/dialects/base_sqlite.ts
@@ -91,16 +91,32 @@ export abstract class BaseSqliteDialect implements DialectContract {
    * Drop all tables inside the database
    */
   async dropAllTables() {
-    await this.client.rawQuery('PRAGMA writable_schema = 1;')
-    await this.client
-      .knexQuery()
-      .delete()
-      .from('sqlite_master')
-      .whereIn('type', ['table', 'index', 'trigger'])
-      .whereNotIn('name', this.config.wipe?.ignoreTables || [])
-
-    await this.client.rawQuery('PRAGMA writable_schema = 0;')
-    await this.client.rawQuery('VACUUM;')
+    const tables = await this.getAllTables()
+
+    /**
+     * Check for foreign key pragma and turn it off if enabled
+     * so that we can drop tables without any issues
+     */
+    const pragma = await this.client.rawQuery('PRAGMA foreign_keys;')
+    if (pragma[0].foreign_keys === 1) {
+      await this.client.rawQuery('PRAGMA foreign_keys = OFF;')
+    }
+
+    /**
+     * Drop all tables
+     */
+    const promises = tables
+      .filter((table) => !this.config.wipe?.ignoreTables?.includes(table))
+      .map((table) => this.client.rawQuery(`DROP TABLE ${table};`))
+
+    await Promise.all(promises)
+
+    /**
+     * Restore foreign key pragma to it's original value
+     */
+    if (pragma[0].foreign_keys === 1) {
+      await this.client.rawQuery('PRAGMA foreign_keys = ON;')
+    }
   }
 
   /**
diff --git a/test-helpers/index.ts b/test-helpers/index.ts
index 0ad47835..53d4ce0a 100644
--- a/test-helpers/index.ts
+++ b/test-helpers/index.ts
@@ -74,12 +74,6 @@ export function getConfig(): ConnectionConfig {
         },
         useNullAsDefault: true,
         debug: !!process.env.DEBUG,
-        pool: {
-          afterCreate(connection, done) {
-            connection.unsafeMode(true)
-            done()
-          },
-        },
       }
     case 'mysql':
       return {
diff --git a/test/database/drop_tables.spec.ts b/test/database/drop_tables.spec.ts
index 97ee6c3f..d863e354 100644
--- a/test/database/drop_tables.spec.ts
+++ b/test/database/drop_tables.spec.ts
@@ -107,4 +107,59 @@ test.group('Query client | drop tables', (group) => {
 
     await connection.disconnect()
   })
+
+  test('drop tables with foreign keys constraint', async ({ assert }) => {
+    const connection = new Connection('primary', getConfig(), logger)
+    connection.connect()
+
+    await connection.client!.schema.createTableIfNotExists('temp_users', (table) => {
+      table.increments('id')
+    })
+
+    await connection.client!.schema.createTableIfNotExists('temp_posts', (table) => {
+      table.increments('id')
+      table.integer('temp_users_id').unsigned().references('id').inTable('temp_users')
+    })
+
+    await connection.client?.table('temp_users').insert({})
+    const user = await connection.client?.table('temp_users').select('id').first()
+    await connection.client?.table('temp_posts').insert({ temp_users_id: user!.id })
+
+    const client = new QueryClient('dual', connection, createEmitter())
+    await client.dialect.dropAllTables(['public'])
+
+    assert.isFalse(await connection.client!.schema.hasTable('temp_users'))
+    assert.isFalse(await connection.client!.schema.hasTable('temp_posts'))
+
+    await connection.disconnect()
+  })
+
+  if (['better-sqlite', 'sqlite'].includes(process.env.DB!)) {
+    test('drop tables when PRAGMA foreign_keys is enabled', async ({ assert }) => {
+      const connection = new Connection('primary', getConfig(), logger)
+      connection.connect()
+
+      await connection.client!.schema.createTable('temp_posts', (table) => {
+        table.increments('id')
+      })
+
+      await connection.client!.schema.createTableIfNotExists('temp_users', (table) => {
+        table.increments('id')
+        table.integer('temp_posts_id').unsigned().references('id').inTable('temp_posts')
+      })
+
+      await connection.client?.table('temp_posts').insert({ id: 1 })
+      await connection.client?.table('temp_users').insert({ id: 1, temp_posts_id: 1 })
+
+      const client = new QueryClient('dual', connection, createEmitter())
+
+      await client.rawQuery('PRAGMA foreign_keys = ON;')
+      await client.dialect.dropAllTables(['public'])
+
+      assert.isFalse(await connection.client!.schema.hasTable('temp_users'))
+      assert.isFalse(await connection.client!.schema.hasTable('temp_posts'))
+
+      await connection.disconnect()
+    })
+  }
 })