drift3: Simplify table extensions

Author: simolus3

future/drift3/lib/src/database/batch.dart

@@ -85,10 +85,7 @@ final class Batch {
     Insertable<Row> row, {
     UpsertClause<Row, RS>? onConflict,
   }) {
-    final stmt = InsertStatement<Row, RS, DatabaseConnectionUser>(
-      _database,
-      table,
-    )..values(row);
+    final stmt = InsertStatement<Row, RS>(_database, table)..values(row);
     if (onConflict != null) {
       stmt.onConflict(onConflict);
     }
@@ -116,10 +113,8 @@ final class Batch {
     required Map<TableColumn, Expression> columns,
     UpsertClause<Row, RS>? onConflict,
   }) {
-    final stmt = InsertStatement<Row, RS, DatabaseConnectionUser>(
-      _database,
-      table,
-    )..fromSelect(select, columns: columns);
+    final stmt = InsertStatement<Row, RS>(_database, table)
+      ..fromSelect(select, columns: columns);
     if (onConflict != null) {
       stmt.onConflict(onConflict);
     }

future/drift3/lib/src/database/connection_user.dart

@@ -457,11 +457,11 @@ abstract base class DatabaseConnectionUser {
 
   /// Starts an [InsertStatement] for a given table. You can use that statement
   /// to write data into the [table] by using [InsertStatement.insert].
-  InsertStatement<Row, RS, DatabaseConnectionUser> into<
+  InsertStatement<Row, RS> into<
     Row extends Object,
     RS extends GeneratedTable<Row, RS>
   >(GeneratedTable<Row, RS> table) {
-    return InsertStatement<Row, RS, DatabaseConnectionUser>(this, table);
+    return InsertStatement<Row, RS>(this, table);
   }
 
   /// Starts an [UpdateStatement] for the given table. You can use that

future/drift3/lib/src/query_builder/compiler.dart

@@ -504,7 +504,6 @@ abstract base class StatementCompiler {
     }
 
     if (delete.returning case final returning?) {
-      statement.resultSetStructure = returning.structure;
       statement.space();
       returning.compileWith(this);
     }
@@ -539,7 +538,6 @@ abstract base class StatementCompiler {
     }
 
     if (update.returning case final returning?) {
-      statement.resultSetStructure = returning.structure;
       statement.space();
       returning.compileWith(this);
     }
@@ -628,6 +626,7 @@ abstract base class StatementCompiler {
   void addReturningClause(ReturningClause returning) {
     // We currently only support the `RETURNING *` format without arbitrary
     // columns.
+    statement.resultSetStructure = returning.structure;
     statement.buffer.write('RETURNING ');
     addResultSetExpressions(returning.structure);
   }
@@ -765,9 +764,16 @@ abstract base class StatementCompiler {
       );
     }
 
+    // The subquery expression can reference outer columns as well, so we need
+    // to make table aliases explicit.
+    final hasMultipleTables = statement.hasMultipleTables;
+    statement.hasMultipleTables = true;
+
     statement.buffer.write('(');
     e.statement.compileWith(this);
     statement.buffer.write(')');
+
+    statement.hasMultipleTables = hasMultipleTables;
   }
 
   void addSubquery(Subquery e) {

future/drift3/lib/src/query_builder/on_table.dart

@@ -1,33 +1,84 @@
 import '../connection/connection.dart';
 import '../database/connection_user.dart';
 import '../database/data_class.dart';
+import '../database/selectable.dart';
+import 'expressions/aggregate.dart';
+import 'expressions/expression.dart';
+import 'schema/result_set.dart';
 import 'schema/table.dart';
 import 'statements/insert.dart';
+import 'statements/select.dart';
+
+/// Easily-accessible methods to compose common operations or statements on
+/// tables or views.
+extension type TableOrViewStatements<
+  Row extends Object,
+  RS extends ResultSet<Row, RS>
+>._((RS, DatabaseConnectionUser) _resultSetWithDatabase) {
+  /// Selects all rows that are in this table.
+  ///
+  /// The returned [Selectable] can be run once as a future with [Selectable.get]
+  /// or as an auto-updating stream with [Selectable.watch].
+  Selectable<Row> all() {
+    return select();
+  }
+
+  /// Counts the rows in this table.
+  ///
+  /// The optional [where] clause can be used to only count rows matching the
+  /// condition, similar to [SimpleSelectStatement.where].
+  ///
+  /// The returned [Selectable] can be run once with [Selectable.getSingle] to
+  /// get the count once, or be watched as a stream with [Selectable.watchSingle].
+  Selectable<int> count({Expression<bool> Function(RS row)? where}) {
+    final count = countAll();
+    final stmt = selectOnly()..addColumns([count]);
+    if (where != null) {
+      stmt.where(where(_resultSetWithDatabase.$1));
+    }
+
+    return stmt.map((row) => row.read(count)!);
+  }
+
+  /// Composes a `SELECT` statement on the captured table or view.
+  ///
+  /// This is equivalent to calling [DatabaseConnectionUser.select].
+  SingleTableSelectStatement<Row, RS> select({bool distinct = false}) {
+    final (rs, db) = _resultSetWithDatabase;
+    return db.select(rs, distinct: distinct);
+  }
+
+  /// Composes a `SELECT` statement only selecting a subset of columns.
+  ///
+  /// This is equivalent to calling [DatabaseConnectionUser.selectOnly].
+  SelectStatement selectOnly({bool distinct = false}) {
+    final (rs, db) = _resultSetWithDatabase;
+    return db.selectOnly(rs, distinct: distinct);
+  }
+}
 
 /// Easily-accessible methods to compose common operations or statements on
 /// tables.
-extension TableStatements<
+extension type TableStatements<
   Row extends Object,
   RS extends GeneratedTable<Row, RS>
->
-    on GeneratedTable<Row, RS> {
+>._(TableOrViewStatements<Row, RS> _super)
+    implements TableOrViewStatements<Row, RS> {
   /// Creates an insert statment to be used to compose an insert on the table.
-  ///
-  /// To later run the statement, first call
-  /// [InsertStatementWithoutDatabase.onDatabase] followed by methods like
-  /// [InsertStatementWithDatabase.insert].
-  InsertStatement<Row, RS, Null> insert() => InsertStatement(null, this);
+  InsertStatement<Row, RS> insert() {
+    final (rs, db) = _resultSetWithDatabase;
+    return InsertStatement(db, rs);
+  }
 
   /// Inserts one row into this table.
   ///
-  /// This is equivalent to calling [InsertStatementWithDatabase.insert] - see
-  /// that method for more information.
+  /// This is equivalent to calling [InsertStatement.insert], see that method
+  /// for more information.
   Future<int> insertOne(
     Insertable<Row> row, {
-    required DatabaseConnectionUser database,
     UpsertClause<Row, RS>? onConflict,
   }) {
-    return insert().onDatabase(database).insert(row, onConflict: onConflict);
+    return insert().insert(row, onConflict: onConflict);
   }
 
   /// Atomically inserts all [rows] into the table.
@@ -37,14 +88,15 @@ extension TableStatements<
   /// the [rows] are in doesn't matter if there are foreign keys between them.
   Future<void> insertAll(
     Iterable<Insertable<Row>> rows, {
-    required DatabaseConnectionUser database,
     UpsertClause<Row, RS>? onConflict,
   }) {
-    return database.transaction(
+    final (rs, db) = _resultSetWithDatabase;
+
+    return db.transaction(
       options: const TransactionOptions(deferForeignKeys: true),
       () async {
-        await database.batch((b) {
-          b.insertAll(this, rows, onConflict: onConflict);
+        await db.batch((b) {
+          b.insertAll(rs, rows, onConflict: onConflict);
         });
       },
     );
@@ -54,16 +106,13 @@ extension TableStatements<
   /// exists already.
   ///
   /// Please note that this method is only available on recent sqlite3 versions.
-  /// See also [InsertStatementWithDatabase.insertOnConflictUpdate].
+  /// See also [InsertStatement.insertOnConflictUpdate].
   /// By default, only the primary key is used for detect uniqueness violations.
   /// If you have further uniqueness constraints, please use the general
   /// [insertOne] method with a [DoUpdate] including those columns in its
   /// [DoUpdate.target].
-  Future<int> insertOnConflictUpdate(
-    Insertable<Row> row, {
-    required DatabaseConnectionUser database,
-  }) {
-    return insert().onDatabase(database).insertOnConflictUpdate(row);
+  Future<int> insertOnConflictUpdate(Insertable<Row> row) {
+    return insert().insertOnConflictUpdate(row);
   }
 
   /// Inserts one row into this table and returns it, along with auto-generated
@@ -75,12 +124,9 @@ extension TableStatements<
   /// use [insertReturningOrNull] instead.
   Future<Row> insertReturning(
     Insertable<Row> row, {
-    required DatabaseConnectionUser database,
     UpsertClause<Row, RS>? onConflict,
   }) {
-    return insert()
-        .onDatabase(database)
-        .insertReturning(row, onConflict: onConflict);
+    return insert().insertReturning(row, onConflict: onConflict);
   }
 
   /// Inserts one row into this table and returns it, along with auto-generated
@@ -90,11 +136,36 @@ extension TableStatements<
   /// [onConflict] clause with a `where` clause was used), returns `null`.
   Future<Row?> insertReturningOrNull(
     Insertable<Row> row, {
-    required DatabaseConnectionUser database,
     UpsertClause<Row, RS>? onConflict,
   }) {
-    return insert()
-        .onDatabase(database)
-        .insertReturningOrNull(row, onConflict: onConflict);
+    return insert().insertReturningOrNull(row, onConflict: onConflict);
+  }
+}
+
+/// Extension providing the [statements] method for tables and views.
+extension GetTableOrViewStatements<
+  Row extends Object,
+  RS extends ResultSet<Row, RS>
+>
+    on ResultSet<Row, RS> {
+  /// Creates a [TableOrViewStatements] instance that can be used to create
+  /// common select statements on this table or view.
+  TableOrViewStatements<Row, RS> statements(DatabaseConnectionUser db) {
+    return TableOrViewStatements._((asSelfType(), db));
+  }
+}
+
+/// Extension providing the [statements] method for tables specifically.
+extension GetTabletatements<
+  Row extends Object,
+  RS extends GeneratedTable<Row, RS>
+>
+    on GeneratedTable<Row, RS> {
+  /// Returns a [TableStatements] instance that can be used to create common
+  /// select, update, insert and delete statements.
+  TableStatements<Row, RS> statements(DatabaseConnectionUser db) {
+    return TableStatements._(
+      GetTableOrViewStatements<Row, RS>(this).statements(db),
+    );
   }
 }

future/drift3/lib/src/query_builder/results.dart

@@ -215,7 +215,7 @@ final class DriftRow {
   /// For result sets that might be absent from some rows (e.g. those added as
   /// outer joins), consider using [readTableOrNull] instead.
   Row readTable<Row extends Object, RS extends ResultSet<Row, RS>>(
-    RS resultSet,
+    ResultSet<Row, RS> resultSet,
   ) {
     final parsed = readTableOrNull<Row, RS>(resultSet);
     if (parsed == null) {
@@ -235,7 +235,7 @@ final class DriftRow {
   /// This does not allow reading result sets whose columns don't exist in this
   /// row at all (e.g. a result set not joined into a select statement).
   Row? readTableOrNull<Row extends Object, RS extends ResultSet<Row, RS>>(
-    RS resultSet,
+    ResultSet<Row, RS> resultSet,
   ) {
     return this.resultSet._mapperFor(resultSet)(raw) as Row?;
   }

future/drift3/lib/src/query_builder/statements/insert.dart

@@ -19,15 +19,14 @@ import 'statement.dart';
 /// An `INSERT` statement in SQL.
 final class InsertStatement<
   Row extends Object,
-  RS extends GeneratedTable<Row, RS>,
-  DB extends DatabaseConnectionUser?
+  RS extends GeneratedTable<Row, RS>
 >
     extends SqlStatement {
   /// The table we're inserting into.
   final GeneratedTable<Row, RS> table;
 
   /// The database used to run this insert statement.
-  final DB _database;
+  final DatabaseConnectionUser _database;
 
   /// The values to use for the rows being inserted.
   InsertSource? source;
@@ -54,7 +53,7 @@ final class InsertStatement<
   ///
   /// Typically, [row] is an instance of the companion class drift generates for
   /// tables, which allows specifying which columns to use.
-  InsertStatement<Row, RS, DB> values(Insertable<Row> row) {
+  InsertStatement<Row, RS> values(Insertable<Row> row) {
     _checkNoSource();
 
     final rawValues = row.toColumns(true);
@@ -101,7 +100,7 @@ final class InsertStatement<
   /// target column, and values are expressions added to the select statement.
   ///
   /// For an example, see the [documentation website](https://drift.simonbinder.eu/docs/advanced-features/joins/#using-selects-as-insert)
-  InsertStatement<Row, RS, DB> fromSelect(
+  InsertStatement<Row, RS> fromSelect(
     BaseSelectStatement select, {
     required Map<TableColumn, Expression> columns,
   }) {
@@ -128,25 +127,12 @@ final class InsertStatement<
   }
 
   /// Adds an [upsert] clause to this insert statement.
-  InsertStatement<Row, RS, DB> onConflict(UpsertClause<Row, RS> upsert) {
+  InsertStatement<Row, RS> onConflict(UpsertClause<Row, RS> upsert) {
     assert(upsertClause == null, 'upsert clause already set');
     upsertClause = upsert;
     return this;
   }
 
-  @override
-  void compileWith(StatementCompiler compiler) {
-    compiler.addInsertStatement(this);
-  }
-}
-
-///
-extension InsertStatementWithDatabase<
-  Row extends Object,
-  RS extends GeneratedTable<Row, RS>,
-  DB extends DatabaseConnectionUser
->
-    on InsertStatement<Row, RS, DB> {
   Future<QueryResult> _run() async {
     final result = await _database.runStatement(this);
     if (result.affectedRows case final rows? when rows > 0) {
@@ -327,23 +313,10 @@ extension InsertStatementWithDatabase<
   Future<int> insertOnConflictUpdate(Insertable<Row> entity) {
     return insert(entity, onConflict: DoUpdate((_) => entity));
   }
-}
 
-/// Exposes the [onDatabase] method to add a database to this insert statement.
-extension InsertStatementWithoutDatabase<
-  Row extends Object,
-  RS extends GeneratedTable<Row, RS>
->
-    on InsertStatement<Row, RS, Null> {
-  /// Attaches a [DatabaseConnectionUser] to this [InsertStatement] so that it
-  /// can be executed.
-  InsertStatement<Row, RS, DatabaseConnectionUser> onDatabase(
-    DatabaseConnectionUser db,
-  ) {
-    return InsertStatement(db, table)
-      ..source = source
-      ..upsertClause = upsertClause
-      ..returning = returning;
+  @override
+  void compileWith(StatementCompiler compiler) {
+    compiler.addInsertStatement(this);
   }
 }