Flag in the close method to abort running queries (#396)

Author: davidmartos96

lib/postgres.dart

@@ -209,7 +209,9 @@ abstract class SessionExecutor {
 
   /// Closes this session, cleaning up resources and forbiding further calls to
   /// [prepare] and [execute].
-  Future<void> close();
+  /// If [force] is set to true, the session will be closed immediately, instead
+  /// of waiting for any pending queries to finish.
+  Future<void> close({bool force = false});
 }
 
 abstract class Connection implements Session, SessionExecutor {

lib/src/pool/pool_impl.dart

@@ -41,7 +41,8 @@ class PoolImplementation<L> implements Pool<L> {
   Future<void> get closed => _semaphore.done;
 
   @override
-  Future<void> close() async {
+  Future<void> close({bool force = false}) async {
+    // TODO: Implement force close.
     await _semaphore.close();
 
     // Connections are closed when they are returned to the pool if it's closed.
@@ -277,9 +278,11 @@ class _PoolConnection implements Connection {
   }
 
   @override
-  Future<void> close() async {
+  Future<void> close({bool force = false}) async {
     // Don't forward the close call, the underlying connection should be re-used
     // when another pool connection is requested.
+
+    // TODO: Implement force close.
   }
 
   @override

lib/src/v3/connection.dart

@@ -582,8 +582,8 @@ class PgConnectionImplementation extends _PgSessionBase implements Connection {
   }
 
   @override
-  Future<void> close() async {
-    await _close(false, null);
+  Future<void> close({bool force = false}) async {
+    await _close(force, null);
   }
 
   Future<void> _close(bool interruptRunning, PgException? cause,

test/pool_test.dart

@@ -209,4 +209,61 @@ void main() {
       expect(name, 'myapp');
     });
   });
+
+  group(skip: 'not implemented', 'force close', () {
+    Future<Pool> openPool(PostgresServer server) async {
+      final pool = Pool.withEndpoints(
+        [await server.endpoint()],
+        settings: PoolSettings(maxConnectionCount: 1),
+      );
+      addTearDown(pool.close);
+      return pool;
+    }
+
+    Future<void> expectPoolClosesForcefully(Pool pool) async {
+      await pool
+          .close(force: true) //
+          // If close takes too long, the test will fail (force=true would not be working correctly)
+          // as it would be waiting for the query to finish
+          .timeout(Duration(seconds: 1));
+      expect(pool.isOpen, isFalse);
+    }
+
+    Future<void> runLongQuery(Session session) {
+      return session.execute('select pg_sleep(10) from pg_stat_activity;');
+    }
+
+    withPostgresServer('pool session', (server) {
+      test('', () async {
+        final pool = await openPool(server);
+        // ignore: unawaited_futures
+        runLongQuery(pool);
+        // let it start
+        await Future.delayed(const Duration(milliseconds: 100));
+        await expectPoolClosesForcefully(pool);
+      });
+    });
+
+    withPostgresServer('tx session', (server) {
+      test('', () async {
+        final pool = await openPool(server);
+        // Ignore async error, it will fail when the connection is closed and it tries to do COMMIT
+        pool.runTx(runLongQuery).ignore();
+        // let it start
+        await Future.delayed(const Duration(milliseconds: 100));
+        await expectPoolClosesForcefully(pool);
+      });
+    });
+
+    withPostgresServer('run session', (server) {
+      test('', () async {
+        final pool = await openPool(server);
+        // ignore: unawaited_futures
+        pool.run(runLongQuery);
+        // let it start
+        await Future.delayed(const Duration(milliseconds: 100));
+        await expectPoolClosesForcefully(pool);
+      });
+    });
+  });
 }

test/v3_close_test.dart

@@ -8,16 +8,23 @@ void main() {
     late Connection conn1;
     late Connection conn2;
 
+    const conn1Name = 'conn1';
+    const conn2Name = 'conn2';
+
     setUp(() async {
       conn1 = await Connection.open(
         await server.endpoint(),
         settings: ConnectionSettings(
-            //transformer: _loggingTransformer('c1'),
-            ),
+          applicationName: conn1Name,
+          //transformer: _loggingTransformer('c1'),
+        ),
       );
 
       conn2 = await Connection.open(
         await server.endpoint(),
+        settings: ConnectionSettings(
+          applicationName: conn2Name,
+        ),
       );
     });
 
@@ -30,9 +37,8 @@ void main() {
       test(
         'with concurrent query: $concurrentQuery',
         () async {
-          final endpoint = await server.endpoint();
           final res = await conn2.execute(
-              "SELECT pid FROM pg_stat_activity where usename = '${endpoint.username}';");
+              "SELECT pid FROM pg_stat_activity where application_name = '$conn1Name';");
           final conn1PID = res.first.first as int;
 
           // Simulate issue by terminating a connection during a query
@@ -49,10 +55,9 @@ void main() {
     }
 
     test('with simple query protocol', () async {
-      final endpoint = await server.endpoint();
       // Get the PID for conn1
       final res = await conn2.execute(
-          "SELECT pid FROM pg_stat_activity where usename = '${endpoint.username}';");
+          "SELECT pid FROM pg_stat_activity where application_name = '$conn1Name';");
       final conn1PID = res.first.first as int;
 
       // ignore: unawaited_futures
@@ -65,6 +70,61 @@ void main() {
           'select pg_terminate_backend($conn1PID) from pg_stat_activity;');
     });
   });
+
+  group('force close', () {
+    Future<Connection> openConnection(PostgresServer server) async {
+      final conn = await Connection.open(await server.endpoint());
+      addTearDown(conn.close);
+      return conn;
+    }
+
+    Future<void> expectConn1ClosesForcefully(Connection conn) async {
+      await conn
+          .close(force: true) //
+          // If close takes too long, the test will fail (force=true would not be working correctly)
+          // as it would be waiting for the query to finish
+          .timeout(Duration(seconds: 1));
+      expect(conn.isOpen, isFalse);
+    }
+
+    Future<void> runLongQuery(Session session) {
+      return session.execute('select pg_sleep(10) from pg_stat_activity;');
+    }
+
+    withPostgresServer('connection session', (server) {
+      test('', () async {
+        final conn = await openConnection(server);
+        // ignore: unawaited_futures
+        runLongQuery(conn);
+        // let it start
+        await Future.delayed(const Duration(milliseconds: 100));
+        await expectConn1ClosesForcefully(conn);
+      });
+    });
+
+    withPostgresServer('tx session', (server) {
+      test('', () async {
+        final conn = await openConnection(server);
+        // ignore: unawaited_futures
+        // Ignore async error, it will fail when the connection is closed and it tries to do COMMIT
+        conn.runTx(runLongQuery).ignore();
+        // let it start
+        await Future.delayed(const Duration(milliseconds: 100));
+        await expectConn1ClosesForcefully(conn);
+      });
+    });
+
+    withPostgresServer('run session', (server) {
+      test('', () async {
+        final conn = await openConnection(server);
+        // ignore: unawaited_futures
+        conn.run(runLongQuery);
+        // let it start
+        await Future.delayed(const Duration(milliseconds: 100));
+        await expectConn1ClosesForcefully(conn);
+      });
+    });
+  });
 }
 
 final _isPostgresException = isA<PgException>();