fix(workspace): force-close connection upon query TimeoutException (#333)

lib/core/database/mysql_connection.dart

@@ -270,7 +270,12 @@ class MysqlConnection {
     if (!isConnected || _conn == null) {
       throw StateError('Not connected to MySQL');
     }
-    return _conn!.execute(sql, params, iterable);
+    try {
+      return await _conn!.execute(sql, params, iterable);
+    } on TimeoutException {
+      unawaited(forceClose());
+      rethrow;
+    }
   }
 
   /// Runs [execute] with an application-level [timeout] (driver limits still apply).
@@ -282,7 +287,12 @@ class MysqlConnection {
   }) async {
     final f = execute(sql, params, iterable);
     if (timeout == null) return f;
-    return f.timeout(timeout);
+    try {
+      return await f.timeout(timeout);
+    } on TimeoutException {
+      unawaited(forceClose());
+      rethrow;
+    }
   }
 
   /// Lists user-visible databases (excludes typical system schemas).

lib/core/database/postgres_connection.dart

@@ -1,3 +1,4 @@
+import 'dart:async';
 import 'dart:io' show SecurityContext;
 
 import 'package:flutter/foundation.dart';
@@ -252,7 +253,27 @@ class PostgresConnection {
     if (!isConnected || _conn == null) {
       throw StateError('Not connected to PostgreSQL');
     }
-    return _conn!.execute(sql, timeout: timeout);
+    try {
+      return await _conn!.execute(sql, timeout: timeout);
+    } on TimeoutException {
+      unawaited(forceClose());
+      rethrow;
+    }
+  }
+
+  /// Runs [execute] with an application-level [timeout] (in addition to driver timeout).
+  Future<Result> executeWithTimeout(
+    String sql, {
+    Duration? timeout,
+  }) async {
+    final f = execute(sql, timeout: timeout);
+    if (timeout == null) return f;
+    try {
+      return await f.timeout(timeout);
+    } on TimeoutException {
+      unawaited(forceClose());
+      rethrow;
+    }
   }
 
   /// Whether the session has an open transaction (PostgreSQL 13+).

lib/core/database/sqlite_connection.dart

@@ -114,11 +114,32 @@ class SqliteConnection {
       throw StateError('Database connection is read-only');
     }
 
-    if (isReadOnlyQuery || hasReturning) {
-      return await _db!.rawQuery(sql, arguments);
-    } else {
-      await _db!.execute(sql, arguments);
-      return [];
+    try {
+      if (isReadOnlyQuery || hasReturning) {
+        return await _db!.rawQuery(sql, arguments);
+      } else {
+        await _db!.execute(sql, arguments);
+        return [];
+      }
+    } on TimeoutException {
+      unawaited(forceClose());
+      rethrow;
+    }
+  }
+
+  /// Runs [execute] with an application-level [timeout].
+  Future<List<Map<String, Object?>>> executeWithTimeout(
+    String sql, {
+    Duration? timeout,
+    List<Object?>? arguments,
+  }) async {
+    final f = execute(sql, arguments);
+    if (timeout == null) return f;
+    try {
+      return await f.timeout(timeout);
+    } on TimeoutException {
+      unawaited(forceClose());
+      rethrow;
     }
   }
 

lib/features/mysql/mysql_sql_workspace.dart

@@ -244,6 +244,7 @@ class _MysqlSqlWorkspaceState extends material.State<MysqlSqlWorkspace> {
         );
       }
     } on TimeoutException catch (e) {
+      unawaited(_lease?.connection.forceClose());
       if (mounted) {
         setState(() {
           _error = e.toString();

lib/features/postgresql/postgres_sql_workspace.dart

@@ -248,6 +248,14 @@ class _PostgresSqlWorkspaceState extends material.State<PostgresSqlWorkspace> {
         _statusLine = 'OK: $cmd';
         _running = false;
       });
+    } on TimeoutException catch (e) {
+      unawaited(_lease?.connection.forceClose());
+      if (mounted) {
+        setState(() {
+          _error = 'Query timed out: ${e.message ?? e}';
+          _running = false;
+        });
+      }
     } on pg.ServerException catch (e) {
       if (mounted) {
         setState(() {
@@ -375,6 +383,14 @@ class _PostgresSqlWorkspaceState extends material.State<PostgresSqlWorkspace> {
           ),
         );
       }
+    } on TimeoutException catch (e) {
+      unawaited(_lease?.connection.forceClose());
+      if (mounted) {
+        setState(() {
+          _error = 'Query timed out: ${e.message ?? e}';
+          _running = false;
+        });
+      }
     } on pg.ServerException catch (e) {
       if (mounted) {
         setState(() {

lib/features/sqlite/sqlite_sql_workspace.dart

@@ -200,6 +200,14 @@ class _SqliteSqlWorkspaceState extends material.State<SqliteSqlWorkspace> {
           ),
         );
       }
+    } on TimeoutException catch (e) {
+      unawaited(_lease?.connection.forceClose());
+      if (mounted) {
+        setState(() {
+          _error = 'Query timed out: ${e.message ?? e}';
+          _running = false;
+        });
+      }
     } catch (e) {
       if (mounted) {
         setState(() {

test/core/database/connection_timeout_protocol_test.dart

@@ -0,0 +1,157 @@
+import 'dart:async';
+import 'package:flutter_test/flutter_test.dart';
+import 'package:mysql_client/mysql_client.dart';
+import 'package:querya_desktop/core/database/mysql_connection.dart';
+import 'package:querya_desktop/core/database/postgres_connection.dart';
+import 'package:querya_desktop/core/database/sqlite_connection.dart';
+import 'package:postgres/postgres.dart' as pg;
+
+class FakeSlowMysqlConnection extends MysqlConnection {
+  FakeSlowMysqlConnection({super.id = 1})
+      : super(
+          name: 'fake_slow_mysql',
+          host: 'localhost',
+          port: 3306,
+          database: 'testdb',
+        );
+
+  bool _connected = true;
+  int forceCloseCount = 0;
+
+  @override
+  bool get isConnected => _connected;
+
+  @override
+  Future<IResultSet> execute(
+    String sql, [
+    Map<String, dynamic>? params,
+    bool iterable = false,
+  ]) async {
+    await Future.delayed(const Duration(seconds: 10));
+    throw Exception('should not reach here');
+  }
+
+  @override
+  Future<void> forceClose() async {
+    forceCloseCount++;
+    _connected = false;
+  }
+}
+
+class FakeSlowPostgresConnection extends PostgresConnection {
+  FakeSlowPostgresConnection({super.id = 1})
+      : super(
+          name: 'fake_slow_pg',
+          host: 'localhost',
+          port: 5432,
+          database: 'postgres',
+        );
+
+  bool _connected = true;
+  int forceCloseCount = 0;
+
+  @override
+  bool get isConnected => _connected;
+
+  @override
+  Future<pg.Result> execute(String sql, {Duration? timeout}) async {
+    await Future.delayed(const Duration(seconds: 10));
+    throw Exception('should not reach here');
+  }
+
+  @override
+  Future<void> forceClose() async {
+    forceCloseCount++;
+    _connected = false;
+  }
+}
+
+class FakeSlowSqliteConnection extends SqliteConnection {
+  FakeSlowSqliteConnection({super.id = 1})
+      : super(
+          name: 'fake_slow_sqlite',
+          path: ':memory:',
+        );
+
+  bool _connected = true;
+  int forceCloseCount = 0;
+
+  @override
+  bool get isConnected => _connected;
+
+  @override
+  Future<List<Map<String, Object?>>> execute(
+    String sql, [
+    List<Object?>? arguments,
+  ]) async {
+    await Future.delayed(const Duration(seconds: 10));
+    throw Exception('should not reach here');
+  }
+
+  @override
+  Future<void> forceClose() async {
+    forceCloseCount++;
+    _connected = false;
+  }
+}
+
+void main() {
+  group('Connection Timeout Protocol Protection', () {
+    test('MysqlConnection.executeWithTimeout force-closes on timeout', () async {
+      final conn = FakeSlowMysqlConnection();
+      expect(conn.isConnected, isTrue);
+
+      try {
+        await conn.executeWithTimeout(
+          'SELECT sleep(100)',
+          timeout: const Duration(milliseconds: 20),
+        );
+        fail('Should have thrown TimeoutException');
+      } on TimeoutException {
+        // Expected
+      }
+
+      await Future.delayed(const Duration(milliseconds: 10));
+      expect(conn.forceCloseCount, 1);
+      expect(conn.isConnected, isFalse);
+    });
+
+    test('PostgresConnection.executeWithTimeout force-closes on TimeoutException', () async {
+      final conn = FakeSlowPostgresConnection();
+      expect(conn.isConnected, isTrue);
+
+      try {
+        await conn.executeWithTimeout(
+          'SELECT pg_sleep(100)',
+          timeout: const Duration(milliseconds: 20),
+        );
+        fail('Should have thrown TimeoutException');
+      } on TimeoutException {
+        // Expected
+      }
+
+      await Future.delayed(const Duration(milliseconds: 10));
+      expect(conn.forceCloseCount, 1);
+      expect(conn.isConnected, isFalse);
+    });
+
+    test('SqliteConnection.executeWithTimeout force-closes on TimeoutException', () async {
+      final conn = FakeSlowSqliteConnection();
+      expect(conn.isConnected, isTrue);
+
+      try {
+        await conn.executeWithTimeout(
+          'SELECT 1',
+          timeout: const Duration(milliseconds: 20),
+        );
+        fail('Should have thrown TimeoutException');
+      } on TimeoutException {
+        // Expected
+      }
+
+      await Future.delayed(const Duration(milliseconds: 10));
+      expect(conn.forceCloseCount, 1);
+      expect(conn.isConnected, isFalse);
+    });
+  });
+}