diff --git a/lib/core/database/mysql_connection.dart b/lib/core/database/mysql_connection.dart index ce6ee2ba..df453aeb 100644 --- a/lib/core/database/mysql_connection.dart +++ b/lib/core/database/mysql_connection.dart @@ -239,11 +239,12 @@ class MysqlConnection { Future execute( String sql, [ Map? params, + bool iterable = false, ]) async { if (!isConnected || _conn == null) { throw StateError('Not connected to MySQL'); } - return _conn!.execute(sql, params); + return _conn!.execute(sql, params, iterable); } /// Runs [execute] with an application-level [timeout] (driver limits still apply). @@ -251,8 +252,9 @@ class MysqlConnection { String sql, { Duration? timeout, Map? params, + bool iterable = false, }) async { - final f = execute(sql, params); + final f = execute(sql, params, iterable); if (timeout == null) return f; return f.timeout(timeout); } diff --git a/lib/core/database/postgres_sql.dart b/lib/core/database/postgres_sql.dart index aec9d344..5cc2dac9 100644 --- a/lib/core/database/postgres_sql.dart +++ b/lib/core/database/postgres_sql.dart @@ -36,3 +36,41 @@ bool shouldSkipImplicitBegin(String sql) { return false; } + +/// Injects a `LIMIT` clause to a read-only query (SELECT, WITH, VALUES) +/// if it does not already contain a `LIMIT` clause. +String injectSqlLimit(String sql, int limit) { + final cleanSql = stripLeadingWhitespaceAndLineComments(sql); + final upper = cleanSql.toUpperCase(); + + final isSelect = upper.startsWith('SELECT') || + upper.startsWith('WITH') || + upper.startsWith('VALUES'); + + if (!isSelect) { + return sql; + } + + // Check if it already has a LIMIT clause + final hasLimit = RegExp(r'\bLIMIT\b', caseSensitive: false).hasMatch(sql); + if (hasLimit) { + return sql; + } + + // Strip trailing whitespace and semicolons to build the body + var body = sql.trimRight(); + var suffix = ''; + + while (true) { + if (body.isEmpty) break; + if (body.endsWith(';')) { + body = body.substring(0, body.length - 1).trimRight(); + suffix = ';$suffix'; + continue; + } + break; + } + + return '$body\nLIMIT $limit$suffix'; +} + diff --git a/lib/features/mysql/mysql_sql_workspace.dart b/lib/features/mysql/mysql_sql_workspace.dart index 580c93b6..2a474351 100644 --- a/lib/features/mysql/mysql_sql_workspace.dart +++ b/lib/features/mysql/mysql_sql_workspace.dart @@ -147,7 +147,7 @@ class _MysqlSqlWorkspaceState extends material.State { } final to = _statementTimeout(); - final rs = await conn.executeWithTimeout(userSql, timeout: to); + final rs = await conn.executeWithTimeout(userSql, timeout: to, iterable: true); if (!mounted) return; @@ -159,8 +159,12 @@ class _MysqlSqlWorkspaceState extends material.State { final rawRows = >[]; var n = 0; final cap = _resultMaxRows; - for (final row in rs.rows) { - if (n >= cap) break; + var truncated = false; + await for (final row in rs.rowsStream) { + if (n >= cap) { + truncated = true; + break; + } rawRows.add( List.generate(row.numOfColumns, (i) => row.colAt(i)), ); @@ -186,11 +190,9 @@ class _MysqlSqlWorkspaceState extends material.State { ? 'OK. Rows affected: $affected.' : 'Command completed.'; } else { - final total = rs.numOfRows; - final truncated = total > cap; _statusLine = truncated - ? 'Showing first $cap of $total row(s).' - : '$total row(s).'; + ? 'Showing first $cap row(s) (result capped).' + : '$n row(s).'; } _running = false; }); diff --git a/lib/features/postgresql/postgres_sql_workspace.dart b/lib/features/postgresql/postgres_sql_workspace.dart index f4c5fd1d..c9158d67 100644 --- a/lib/features/postgresql/postgres_sql_workspace.dart +++ b/lib/features/postgresql/postgres_sql_workspace.dart @@ -261,7 +261,7 @@ class _PostgresSqlWorkspaceState extends material.State { Future _execute() async { final userSql = _sqlController.text.trim(); if (userSql.isEmpty) return; - var sql = userSql; + var sql = injectSqlLimit(userSql, _resultMaxRows); setState(() { _running = true; @@ -323,9 +323,9 @@ class _PostgresSqlWorkspaceState extends material.State { _statusLine = 'Command completed. Rows affected: ${result.affectedRows}.'; } else { - final truncated = result.length > cap; + final truncated = result.length >= cap; _statusLine = truncated - ? 'Showing first $cap of ${result.length} row(s).' + ? 'Showing first $cap row(s) (result capped).' : '${result.length} row(s).'; } _running = false; diff --git a/test/core/database/postgres_sql_test.dart b/test/core/database/postgres_sql_test.dart index 6b7c70cc..0d49cac0 100644 --- a/test/core/database/postgres_sql_test.dart +++ b/test/core/database/postgres_sql_test.dart @@ -79,4 +79,33 @@ void main() { expect(shouldSkipImplicitBegin('UPDATE t SET x = 1'), isFalse); }); }); + + group('injectSqlLimit', () { + test('appends LIMIT to select query without limit', () { + expect(injectSqlLimit('SELECT * FROM users', 5000), 'SELECT * FROM users\nLIMIT 5000'); + }); + + test('handles trailing semicolons', () { + expect(injectSqlLimit('SELECT * FROM users;', 5000), 'SELECT * FROM users\nLIMIT 5000;'); + expect(injectSqlLimit('SELECT * FROM users; ', 5000), 'SELECT * FROM users\nLIMIT 5000;'); + expect(injectSqlLimit('SELECT * FROM users;;', 5000), 'SELECT * FROM users\nLIMIT 5000;;'); + }); + + test('does not append LIMIT if LIMIT already exists', () { + expect(injectSqlLimit('SELECT * FROM users LIMIT 10', 5000), 'SELECT * FROM users LIMIT 10'); + expect(injectSqlLimit('SELECT * FROM users limit 10;', 5000), 'SELECT * FROM users limit 10;'); + }); + + test('does not modify non-select/non-read queries', () { + expect(injectSqlLimit('INSERT INTO users VALUES (1)', 5000), 'INSERT INTO users VALUES (1)'); + expect(injectSqlLimit('UPDATE users SET x = 1', 5000), 'UPDATE users SET x = 1'); + }); + + test('appends LIMIT to WITH query', () { + expect( + injectSqlLimit('WITH t AS (SELECT * FROM users) SELECT * FROM t;', 5000), + 'WITH t AS (SELECT * FROM users) SELECT * FROM t\nLIMIT 5000;', + ); + }); + }); }