From 4defed1fd249f7fb2fe1b1124d6df0ea584418c8 Mon Sep 17 00:00:00 2001
From: john <john.babawale@habaripay.com>
Date: Wed, 17 Apr 2024 12:30:24 +0100
Subject: [PATCH] feat: allow db:drop with force for PostgreSQL > v13

---
 src/commands/database.js | 56 ++++++++++++++++++++++++++++++----------
 test/db/db-drop.test.js  | 24 +++++++++++++++++
 2 files changed, 66 insertions(+), 14 deletions(-)

diff --git a/src/commands/database.js b/src/commands/database.js
index 22c2d59a4..c61df9b5e 100644
--- a/src/commands/database.js
+++ b/src/commands/database.js
@@ -29,6 +29,11 @@ exports.builder = (yargs) =>
     .option('template', {
       describe: 'Pass template option to dialect, PostgreSQL only',
       type: 'string',
+    })
+    .option('force', {
+      describe:
+        'Pass force option to dialect with db:drop, PostgreSQL > v13 only',
+      type: 'boolean',
     }).argv;
 
 exports.handler = async function (args) {
@@ -45,18 +50,16 @@ exports.handler = async function (args) {
     'encoding',
     'ctype',
     'template',
+    'force',
   ]);
 
-  const queryInterface = sequelize.getQueryInterface();
-  const queryGenerator =
-    queryInterface.queryGenerator || queryInterface.QueryGenerator;
-
-  const query = getCreateDatabaseQuery(sequelize, config, options);
+  const createQuery = getCreateDatabaseQuery(sequelize, config, options);
+  const dropQuery = getDropDatabaseQuery(sequelize, config, options);
 
   switch (command) {
     case 'db:create':
       await sequelize
-        .query(query, {
+        .query(createQuery, {
           type: sequelize.QueryTypes.RAW,
         })
         .catch((e) => helpers.view.error(e));
@@ -66,14 +69,9 @@ exports.handler = async function (args) {
       break;
     case 'db:drop':
       await sequelize
-        .query(
-          `DROP DATABASE IF EXISTS ${queryGenerator.quoteIdentifier(
-            config.database
-          )}`,
-          {
-            type: sequelize.QueryTypes.RAW,
-          }
-        )
+        .query(dropQuery, {
+          type: sequelize.QueryTypes.RAW,
+        })
         .catch((e) => helpers.view.error(e));
 
       helpers.view.log('Database', clc.blueBright(config.database), 'dropped.');
@@ -145,6 +143,36 @@ function getCreateDatabaseQuery(sequelize, config, options) {
   }
 }
 
+function getDropDatabaseQuery(sequelize, config, options) {
+  const queryInterface = sequelize.getQueryInterface();
+  const queryGenerator =
+    queryInterface.queryGenerator || queryInterface.QueryGenerator;
+
+  switch (config.dialect) {
+    // Adds the force option for WITH(FORCE) to drop a database that has connected users, fallback to default drop if version lower
+    // for postgres v13 and above see manual https://www.postgresql.org/docs/current/sql-dropdatabase.html
+    case 'postgres':
+      if (options.force) {
+        helpers.view.log(
+          clc.redBright(
+            `WARNING :: Dropping database with force for v13 and above only (this will drop regardless of connected users) `
+          )
+        );
+        return `DROP DATABASE IF EXISTS ${queryGenerator.quoteIdentifier(
+          config.database
+        )} WITH (FORCE)`;
+      } else
+        return `DROP DATABASE IF EXISTS ${queryGenerator.quoteIdentifier(
+          config.database
+        )}`;
+
+    default:
+      return `DROP DATABASE IF EXISTS ${queryGenerator.quoteIdentifier(
+        config.database
+      )}`;
+  }
+}
+
 function getDatabaseLessSequelize() {
   let config = null;
 
diff --git a/test/db/db-drop.test.js b/test/db/db-drop.test.js
index 832d02136..f7b785e72 100644
--- a/test/db/db-drop.test.js
+++ b/test/db/db-drop.test.js
@@ -48,6 +48,30 @@ describe(Support.getTestDialectTeaser('db:drop'), () => {
         }
       );
     });
+    it('correctly drops database with force', function (done) {
+      const databaseName = `my_test_db_${_.random(10000, 99999)}`;
+      prepare(
+        'db:drop --force',
+        () => {
+          this.sequelize
+            .query(
+              `SELECT 1 as exists FROM pg_database WHERE datname = '${databaseName}';`,
+              {
+                type: this.sequelize.QueryTypes.SELECT,
+              }
+            )
+            .then((result) => {
+              expect(result).to.be.empty;
+              done();
+            });
+        },
+        {
+          config: {
+            database: databaseName,
+          },
+        }
+      );
+    });
   }
 
   if (Support.dialectIsMySQL()) {