Repository: geldata/gel
Pull request: Support closing all connections to a database on database drop
`drop database` will now interrupt and close idle frontend connections before the actual backend database is dropped. Note, idle-in-transaction frontend connections will also be closed.

@msullivan you can configure `close_frontend_conns=True` conditionally in `execute.pyx`.

Fixes #6776 

- [x] Add a test

Problem statement:
Provide a way to force close all connections to a database
This would be useful for certain anticipated branch-related dev workflows.

It doesn't need to be instant, necessarily. Maybe it would be enough to close all the client connections and then wait for everything to clean itself up?

Hints from issue discussion:
I am assigning this to @fantix aspirationally; if you don't think you'll have time, we can sync up about what I would need to do to implement it.
(And if you do have some time, you could make `DROP DATABASE` apply this unconditionally and then leave me to worry about making it conditional in some useful way.)
I was about to post a question on Discord of this feature (request). Would be very useful!

Commit message: make it work on the active queries

Patch to review:
diff --git a/edb/edgeql-parser/src/keywords.rs b/edb/edgeql-parser/src/keywords.rs
index 9d7e438a7fb..3b994b22825 100644
--- a/edb/edgeql-parser/src/keywords.rs
+++ b/edb/edgeql-parser/src/keywords.rs
@@ -20,6 +20,7 @@ pub const UNRESERVED_KEYWORDS: phf::Set<&str> = phf_set!(
     "committed",
     "config",
     "conflict",
+    "connections",
     "constraint",
     "cube",
     "current",
diff --git a/edb/edgeql/ast.py b/edb/edgeql/ast.py
index a936f104d59..3cab1f853c1 100644
--- a/edb/edgeql/ast.py
+++ b/edb/edgeql/ast.py
@@ -883,11 +883,11 @@ class CreateDatabase(CreateObject, DatabaseCommand):
 
 
 class AlterDatabase(AlterObject, DatabaseCommand):
-    pass
+    reset_connections: bool = False
 
 
 class DropDatabase(DropObject, DatabaseCommand):
-    pass
+    reset_connections: bool = False
 
 
 class ExtensionPackageCommand(GlobalObjectCommand):
diff --git a/edb/edgeql/parser/grammar/ddl.py b/edb/edgeql/parser/grammar/ddl.py
index 8dd52011c6a..12ad657587d 100644
--- a/edb/edgeql/parser/grammar/ddl.py
+++ b/edb/edgeql/parser/grammar/ddl.py
@@ -851,7 +851,12 @@ def reduce_create_data_branch(self, *kids):
 #
 class DropBranchStmt(Nonterm):
     def reduce_DROP_BRANCH_DatabaseName(self, *kids):
-        self.val = qlast.DropDatabase(name=kids[2].val)
+        _, _, name = kids
+        self.val = qlast.DropDatabase(name=name.val)
+
+    def reduce_DROP_BRANCH_RESET_CONNECTIONS_DatabaseName(self, *kids):
+        _, _, _, _, name = kids
+        self.val = qlast.DropDatabase(name=name.val, reset_connections=True)
 
 
 #
@@ -867,15 +872,24 @@ def reduce_DROP_BRANCH_DatabaseName(self, *kids):
 
 
 class AlterBranchStmt(Nonterm):
-    def reduce_ALTER_BRANCH_DatabaseName_AlterBranchCommandsBlock(
-        self, *kids
-    ):
+    def reduce_ALTER_BRANCH_DatabaseName_AlterBranchCommandsBlock(self, *kids):
         _, _, name, commands = kids
         self.val = qlast.AlterDatabase(
             name=name.val,
             commands=commands.val,
         )
 
+    def reduce_alter_branch_reset_connections(self, *kids):
+        """%reduce
+            ALTER BRANCH RESET CONNECTIONS DatabaseName AlterBranchCommandsBlock
+        """
+        _, _, _, _, name, commands = kids
+        self.val = qlast.AlterDatabase(
+            name=name.val,
+            commands=commands.val,
+            reset_connections=True,
+        )
+
 
 #
 # EXTENSION PACKAGE
diff --git a/edb/server/compiler/compiler.py b/edb/server/compiler/compiler.py
index 8fff7daf5ea..709f86c7406 100644
--- a/edb/server/compiler/compiler.py
+++ b/edb/server/compiler/compiler.py
@@ -2452,6 +2452,7 @@ def _try_compile(
             unit.sql = comp.sql
             unit.create_db = comp.create_db
             unit.drop_db = comp.drop_db
+            unit.drop_db_reset_connections = comp.drop_db_reset_connections
             unit.create_db_template = comp.create_db_template
             unit.create_db_mode = comp.create_db_mode
             unit.ddl_stmt_id = comp.ddl_stmt_id
diff --git a/edb/server/compiler/dbstate.py b/edb/server/compiler/dbstate.py
index 6d244e73173..6f3867b884b 100644
--- a/edb/server/compiler/dbstate.py
+++ b/edb/server/compiler/dbstate.py
@@ -147,6 +147,7 @@ class DDLQuery(BaseQuery):
     single_unit: bool = False
     create_db: Optional[str] = None
     drop_db: Optional[str] = None
+    drop_db_reset_connections: bool = False
     create_db_template: Optional[str] = None
     create_db_mode: Optional[qlast.BranchType] = None
     ddl_stmt_id: Optional[str] = None
@@ -268,6 +269,7 @@ class QueryUnit:
     # close all inactive unused pooled connections to it.
     create_db: Optional[str] = None
     drop_db: Optional[str] = None
+    drop_db_reset_connections: bool = False
 
     # If non-None, contains a name of the DB that will be used as
     # a template database to create the database. The server should
diff --git a/edb/server/compiler/ddl.py b/edb/server/compiler/ddl.py
index 5ad36efb460..3a1bca06e66 100644
--- a/edb/server/compiler/ddl.py
+++ b/edb/server/compiler/ddl.py
@@ -248,10 +248,12 @@ def compile_and_apply_ddl_stmt(
 
     create_db = None
     drop_db = None
+    drop_db_reset_connections = False
     create_db_template = None
     create_db_mode = None
     if isinstance(stmt, qlast.DropDatabase):
         drop_db = stmt.name.name
+        drop_db_reset_connections = stmt.reset_connections
     elif isinstance(stmt, qlast.CreateDatabase):
         create_db = stmt.name.name
         create_db_template = stmt.template.name if stmt.template else None
@@ -261,6 +263,7 @@ def compile_and_apply_ddl_stmt(
             if isinstance(cmd, qlast.Rename):
                 drop_db = stmt.name.name
                 create_db = cmd.new_name.name
+                drop_db_reset_connections = stmt.reset_connections
 
     if debug.flags.delta_execute_ddl:
         debug.header('Delta Script (DDL Only)')
@@ -283,6 +286,7 @@ def compile_and_apply_ddl_stmt(
         ),
         create_db=create_db,
         drop_db=drop_db,
+        drop_db_reset_connections=drop_db_reset_connections,
         create_db_template=create_db_template,
         create_db_mode=create_db_mode,
         ddl_stmt_id=ddl_stmt_id,
diff --git a/edb/server/protocol/binary.pyx b/edb/server/protocol/binary.pyx
index 007c4ccf4e0..575f5e1dc0f 100644
--- a/edb/server/protocol/binary.pyx
+++ b/edb/server/protocol/binary.pyx
@@ -1088,7 +1088,7 @@ cdef class EdgeConnection(frontend.FrontendConnection):
         self.write_log(
             EdgeSeverity.EDGE_SEVERITY_NOTICE,
             errors.LogMessage.get_code(),
-            'server is stopped; disconnecting now')
+            'requested to stop; disconnecting now')
 
     async def recover_from_error(self):
         # Consume all messages until sync.
diff --git a/edb/server/protocol/execute.pyx b/edb/server/protocol/execute.pyx
index 5631df1c7e3..f6dcc1f6658 100644
--- a/edb/server/protocol/execute.pyx
+++ b/edb/server/protocol/execute.pyx
@@ -97,7 +97,11 @@ async def execute(
                 dbv.dbname,
             )
         if query_unit.drop_db:
-            await tenant.on_before_drop_db(query_unit.drop_db, dbv.dbname)
+            await tenant.on_before_drop_db(
+                query_unit.drop_db,
+                dbv.dbname,
+                close_frontend_conns=query_unit.drop_db_reset_connections,
+            )
         if query_unit.system_config:
             await execute_system_config(be_conn, dbv, query_unit, state)
         else:
diff --git a/edb/server/protocol/frontend.pxd b/edb/server/protocol/frontend.pxd
index f96b2580214..36c2e2dd0d9 100644
--- a/edb/server/protocol/frontend.pxd
+++ b/edb/server/protocol/frontend.pxd
@@ -35,7 +35,7 @@ cdef class FrontendConnection(AbstractFrontendConnection):
         object server
         readonly object tenant
         object loop
-        str dbname
+        readonly str dbname
         str username
         dbview.Database database
 
diff --git a/edb/server/protocol/frontend.pyx b/edb/server/protocol/frontend.pyx
index c93c7c34605..c3b99900e26 100644
--- a/edb/server/protocol/frontend.pyx
+++ b/edb/server/protocol/frontend.pyx
@@ -454,7 +454,7 @@ cdef class FrontendConnection(AbstractFrontendConnection):
             self._transport.abort()
             self._transport = None
 
-    def stop(self):
+    def request_stop(self):
         # Actively stop a frontend connection - this is used by the server
         # when it's stopping.
         self._stop_requested = True
diff --git a/edb/server/protocol/pg_ext.pyx b/edb/server/protocol/pg_ext.pyx
index f0e910e3220..19c78ebf19f 100644
--- a/edb/server/protocol/pg_ext.pyx
+++ b/edb/server/protocol/pg_ext.pyx
@@ -455,7 +455,7 @@ cdef class PgConnection(frontend.FrontendConnection):
                     if self.debug:
                         self.debug_print("CancelRequest", pid, secret)
                     self.server.cancel_pgext_connection(pid, secret)
-                    self.stop()
+                    self.request_stop()
                     break
 
                 elif proto_ver_minor == 5679:  # SSLRequest
@@ -760,7 +760,7 @@ cdef class PgConnection(frontend.FrontendConnection):
             self.write_error(ex)
             self.write(self.ready_for_query())
             self.flush()
-            self.stop()
+            self.request_stop()
 
     async def _main_step(self, char mtype):
         cdef:
diff --git a/edb/server/server.py b/edb/server/server.py
index 4beff7647ac..38937fac3b3 100644
--- a/edb/server/server.py
+++ b/edb/server/server.py
@@ -25,6 +25,7 @@
 import asyncio
 import collections
 import ipaddress
+import itertools
 import json
 import logging
 import os
@@ -425,6 +426,13 @@ def stmt_cache_size(self) -> int | None:
     def system_compile_cache(self):
         return self._system_compile_cache
 
+    def request_stop_fe_conns(self, dbname: str) -> None:
+        for conn in itertools.chain(
+            self._binary_conns.keys(), self._pgext_conns.values()
+        ):
+            if conn.dbname == dbname:
+                conn.request_stop()
+
     def _idle_gc_collector(self):
         try:
             self._idle_gc_handler = None
@@ -987,11 +995,11 @@ async def stop(self):
         self._servers = {}
 
         for conn in self._binary_conns:
-            conn.stop()
+            conn.request_stop()
         self._binary_conns.clear()
 
         for conn in self._pgext_conns.values():
-            conn.stop()
+            conn.request_stop()
         self._pgext_conns.clear()
 
     async def serve_forever(self):
diff --git a/edb/server/tenant.py b/edb/server/tenant.py
index ee8a44cddb9..f9cc546748e 100644
--- a/edb/server/tenant.py
+++ b/edb/server/tenant.py
@@ -728,34 +728,40 @@ def is_database_connectable(self, dbname: str) -> bool:
             and dbname not in self._block_new_connections
         )
 
-    async def ensure_database_not_connected(self, dbname: str) -> None:
+    async def ensure_database_not_connected(
+        self, dbname: str, close_frontend_conns: bool = False
+    ) -> None:
         if self._dbindex and self._dbindex.count_connections(dbname):
-            # If there are open EdgeDB connections to the `dbname` DB
-            # just raise the error Postgres would have raised itself.
-            raise errors.ExecutionError(
-                f"database branch {dbname!r} is being accessed by other users"
-            )
-        else:
-            self._block_new_connections.add(dbname)
+            if close_frontend_conns:
+                self._server.request_stop_fe_conns(dbname)
+            else:
+                # If there are open EdgeDB connections to the `dbname` DB
+                # just raise the error Postgres would have raised itself.
+                raise errors.ExecutionError(
+                    f"database branch {dbname!r} is being accessed by "
+                    f"other users"
+                )
 
-            # Prune our inactive connections.
-            await self._pg_pool.prune_inactive_connections(dbname)
+        self._block_new_connections.add(dbname)
 
-            # Signal adjacent servers to prune their connections to this
-            # database.
-            await self.signal_sysevent(
-                "ensure-database-not-used", dbname=dbname
-            )
+        # Signal adjacent servers to prune their connections to this
+        # database.
+        await self.signal_sysevent(
+            "ensure-database-not-used", dbname=dbname
+        )
 
-            rloop = retryloop.RetryLoop(
-                timeout=10.0,
-                ignore=errors.ExecutionError,
-            )
+        rloop = retryloop.RetryLoop(
+            timeout=10.0,
+            ignore=errors.ExecutionError,
+        )
+
+        async for iteration in rloop:
+            async with iteration:
+                # Prune our inactive connections.  (Do it in the loop
+                # to help in the close_frontend_conns situation.)
+                await self._pg_pool.prune_inactive_connections(dbname)
 
-            async for iteration in rloop:
-                async with iteration:
-                    # Verify we are disconnected
-                    await self._pg_ensure_database_not_connected(dbname)
+                await self._pg_ensure_database_not_connected(dbname)
 
     async def _pg_ensure_database_not_connected(self, dbname: str) -> None:
         async with self.use_sys_pgcon() as pgcon:
@@ -1163,14 +1169,19 @@ def reload(self):
         self._start_watching_files()
 
     async def on_before_drop_db(
-        self, dbname: str, current_dbname: str
+        self,
+        dbname: str,
+        current_dbname: str,
+        close_frontend_conns: bool = False,
     ) -> None:
         if current_dbname == dbname:
             raise errors.ExecutionError(
                 f"cannot drop the currently open database branch {dbname!r}"
             )
 
-        await self.ensure_database_not_connected(dbname)
+        await self.ensure_database_not_connected(
+            dbname, close_frontend_conns=close_frontend_conns
+        )
 
     async def on_before_create_db_from_template(
         self, dbname: str, current_dbname: str
diff --git a/tests/test_database.py b/tests/test_database.py
index f21b4d1a844..2f5dad4ac75 100644
--- a/tests/test_database.py
+++ b/tests/test_database.py
@@ -17,6 +17,7 @@
 #
 
 
+import asyncio
 import edgedb
 
 from edb.schema import defines as s_def
@@ -238,6 +239,112 @@ async def test_branch_drop_recreate(self):
         finally:
             await tb.drop_db(self.con, 'test_db_drop')
 
+    async def _test_branch_drop_disconnect(
+        self, *, with_transaction, with_query
+    ):
+        if not self.has_create_database:
+            self.skipTest("create branch is not supported by the backend")
+
+        await self.con.execute('CREATE EMPTY BRANCH test_db_disconnect;')
+        conn = await self.connect(database='test_db_disconnect')
+        sleeping = None
+
+        try:
+            if with_transaction:
+                await conn.query('START TRANSACTION')
+
+            dbname = await conn.query(
+                'SELECT sys::get_current_database();')
+            self.assertEqual(dbname, ['test_db_disconnect'])
+
+            if with_query:
+                await conn.query('select sys::_sleep(0)')
+                sleeping = asyncio.create_task(
+                    conn.query('select sys::_sleep(3)')
+                )
+                await asyncio.sleep(1)
+
+            # Drop branch while the frontend connection is active
+            await self.con.execute(
+                'DROP BRANCH RESET CONNECTIONS test_db_disconnect'
+            )
+
+            if with_query:
+                try:
+                    await sleeping
+                except edgedb.EdgeDBError:
+                    pass
+                sleeping = None
+
+            # The frontend connection should be closed by the server now
+            self.assertTrue(conn.is_closed())
+        finally:
+            if sleeping:
+                try:
+                    await sleeping
+                except Exception:
+                    pass
+
+            await conn.aclose()
+            try:
+                await tb.drop_db(self.con, 'test_db_disconnect')
+            except edgedb.UnknownDatabaseError:
+                pass
+
+    async def test_branch_drop_disconnect_01(self):
+        await self._test_branch_drop_disconnect(
+            with_transaction=False, with_query=False,
+        )
+
+    async def test_branch_drop_disconnect_02(self):
+        await self._test_branch_drop_disconnect(
+            with_transaction=True, with_query=False,
+        )
+
+    async def test_branch_drop_disconnect_03(self):
+        await self._test_branch_drop_disconnect(
+            with_transaction=False, with_query=True,
+        )
+
+    async def test_branch_drop_disconnect_04(self):
+        await self._test_branch_drop_disconnect(
+            with_transaction=True, with_query=True,
+        )
+
+    async def test_branch_rename_disconnect(self):
+        if not self.has_create_database:
+            self.skipTest("create branch is not supported by the backend")
+
+        await self.con.execute('CREATE EMPTY BRANCH test_db_rename;')
+        conn = await self.connect(database='test_db_rename')
+
+        try:
+            dbname = await conn.query(
+                'SELECT sys::get_current_database();')
+            self.assertEqual(dbname, ['test_db_rename'])
+
+            # Drop branch while the frontend connection is active
+            await self.con.execute('''
+                ALTER BRANCH RESET CONNECTIONS test_db_rename
+                RENAME TO test_db_rename2
+            ''')
+
+            # The frontend connection should be closed by the server now
+            self.assertTrue(conn.is_closed())
+
+            conn2 = await self.connect(database='test_db_rename2')
+            dbname = await conn2.query(
+                'SELECT sys::get_current_database();')
+            self.assertEqual(dbname, ['test_db_rename2'])
+
+        finally:
+            await conn.aclose()
+            await conn2.aclose()
+            try:
+                await tb.drop_db(self.con, 'test_db_rename')
+            except edgedb.UnknownDatabaseError:
+                await tb.drop_db(self.con, 'test_db_rename2')
+
     async def test_branch_non_exist_template(self):
         if not self.has_create_database:
             self.skipTest("create branch is not supported by the backend")


Code review comments:
diff_hunk: @@ -851,7 +851,12 @@ def reduce_create_data_branch(self, *kids):
 #
 class DropBranchStmt(Nonterm):
     def reduce_DROP_BRANCH_DatabaseName(self, *kids):
-        self.val = qlast.DropDatabase(name=kids[2].val)
+        _, _, name = kids
+        self.val = qlast.DropDatabase(name=name.val)
+
+    def reduce_DROP_BRANCH_RESET_CONNECTIONS_DatabaseName(self, *kids):
path: edb/edgeql/parser/grammar/ddl.py
text: @user1:
Not sure I like the formulation. It doesn't read very well and would scale poorly if we add more options. I actually like PostgreSQL's `drop database foo [with option, ...]` somewhat more. I also like `force` as a shorter spelling (a-la Postgres also)

@user2:
Sure