Add Dart API to build compound statements

Closes #3346

Author: simolus3

docs/docs/dart_api/select.md

@@ -5,9 +5,6 @@ description: Select rows or individual columns from tables in Dart
 
 ---
 
-
-
-
 This page describes how to write `SELECT` statements with drift's dart_api.
 To make examples easier to grasp, they're referencing two common tables forming
 the basis of a todo-list app:
@@ -288,8 +285,6 @@ Any statement can be used as a subquery. But be aware that, unlike [subquery exp
 
 ## JSON support
 
-
-
 sqlite3 has great support for [JSON operators](https://sqlite.org/json1.html) that are also available
 in drift (under the additional `'package:drift/extensions/json1.dart'` import).
 JSON support is helpful when storing a dynamic structure that is best represented with JSON, or when
@@ -333,3 +328,34 @@ at all.
 Instead, the expressions in the list passed to `selectExpressions` are evaluated in a standalone
 select statement and can be parsed from the `TypedResult` class returned when evaluating the
 query.
+
+## Compound selects
+
+With compound selects, the results of multiple selects statements can be returned at once.
+Different operators are available to apply set operations on queries, namely:
+
+1. `UNION ALL` and `UNION`: Returns the results of two select statements in a select, with
+   duplicates included or filtered, respectively.
+2. `EXCEPT`: Returns all rows of the first select statement that did not appear in the second
+   query.
+3. `INTERSECT`: Returns all rows that were returned by both select statements.
+
+As an example, consider the tables used to track todo items introduced [in the article on tables](tables.md#defining-tables). Here, one table stores todo items and another table defines categories
+that can be used to group these items.
+Now, perhaps you want to query how many items are assigned to each category, as well as the amount
+of items not in any category.
+The first query can be written with a `groupBy` on categories and a [subquery](expressions.md#subqueries)
+to count associated todo items.
+When grouping on the categories table though, there will be no "null" group. So, one way to resolve
+everything in a single query is to write another query and use `unionAll`:
+
+{{ load_snippet('compound','lib/snippets/dart_api/select.dart.excerpt.json') }}
+
+This query will return one row for each category, counting associated todo items. Also, it includes
+a final row without a category description reporting the count of todo items outside of categories.
+
+With all of these operators, all involved queries must return compatible rows. This is because
+the queries are ultimately reported as a single result set, so they must return the same column
+types.
+It is possible to apply a `LIMIT` and `ORDER BY` clause to compound select statements, but only
+to the first statement (the one on which `union`, `unionAll`, `except` or `intersect` is called).

docs/lib/snippets/dart_api/select.dart

@@ -211,4 +211,26 @@ extension SelectExamples on CanUseCommonTables {
     return row.read(todoItemExists)!;
   }
   // #enddocregion hasTodoItem
+
+  // #docregion compound
+  Future<List<(String?, int)>> todoItemsInCategory() async {
+    final countWithCategory = subqueryExpression<int>(selectOnly(todoItems)
+      ..addColumns([countAll()])
+      ..where(todoItems.category.equalsExp(categories.id)));
+
+    final countWithoutCategory = subqueryExpression<int>(selectOnly(todoItems)
+      ..addColumns([countAll()])
+      ..where(todoItems.category.isNull()));
+
+    final query = db.selectOnly(categories)
+      ..addColumns([categories.name, countWithoutCategory])
+      ..groupBy([categories.id]);
+    query.unionAll(db.selectExpressions(
+        [const Constant<String>(null), countWithoutCategory]));
+
+    return query
+        .map((row) => (row.read(categories.name), row.read(countWithCategory)!))
+        .get();
+  }
+  // #enddocregion compound
 }

drift/CHANGELOG.md

@@ -1,3 +1,7 @@
+## 2.23.0-dev
+
+- Allow building compound select statements in Dart.
+
 ## 2.22.1
 
 - Fix generated SQL for `insertFromSelect` statements with upserts.

drift/lib/src/runtime/api/connection_user.dart

@@ -268,7 +268,8 @@ abstract class DatabaseConnectionUser {
   /// final row = await selectExpressions([currentDateAndTime]).getSingle();
   /// final databaseTime = row.read(currentDateAndTime)!;
   /// ```
-  Selectable<TypedResult> selectExpressions(Iterable<Expression> columns) {
+  BaseSelectStatement<TypedResult> selectExpressions(
+      Iterable<Expression> columns) {
     return SelectWithoutTables(this, columns);
   }
 

drift/lib/src/runtime/query_builder/statements/query.dart

@@ -37,12 +37,14 @@ abstract class Query<T extends HasResultSet, D> extends Component {
   @visibleForOverriding
   void writeStartPart(GenerationContext ctx);
 
-  @override
-  void writeInto(GenerationContext context) {
+  void _writeInto(GenerationContext context,
+      {bool withOrderByAndLimit = true}) {
     // whether we need to insert a space before writing the next component
     var needsWhitespace = false;
 
-    void writeWithSpace(Component? component) {
+    void writeWithSpace(
+      Component? component,
+    ) {
       if (component == null) return;
 
       if (needsWhitespace) context.writeWhitespace();
@@ -55,8 +57,10 @@ abstract class Query<T extends HasResultSet, D> extends Component {
 
     writeWithSpace(whereExpr);
     writeWithSpace(_groupBy);
-    writeWithSpace(orderByExpr);
-    writeWithSpace(limitExpr);
+    if (withOrderByAndLimit) {
+      writeWithSpace(orderByExpr);
+      writeWithSpace(limitExpr);
+    }
 
     if (writeReturningClause) {
       if (needsWhitespace) context.writeWhitespace();
@@ -65,6 +69,11 @@ abstract class Query<T extends HasResultSet, D> extends Component {
     }
   }
 
+  @override
+  void writeInto(GenerationContext context) {
+    _writeInto(context);
+  }
+
   /// Constructs the query that can then be sent to the database executor.
   ///
   /// This is used internally by drift to run the query. Users should use the

drift/lib/src/runtime/query_builder/statements/select/select.dart

@@ -8,7 +8,7 @@ typedef OrderClauseGenerator<T> = OrderingTerm Function(T tbl);
 ///
 /// Users are not allowed to extend, implement or mix-in this class.
 @sealed
-abstract class BaseSelectStatement<Row> extends Component {
+abstract class BaseSelectStatement<Row> extends Component with Selectable<Row> {
   Iterable<(Expression, String)> get _expandedColumns;
 
   /// The name for the given [expression] in the result set, or `null` if
@@ -192,6 +192,7 @@ final class SelectWithoutTables extends BaseSelectStatement<TypedResult>
 
   @override
   void writeInto(GenerationContext context) {
+    final isRoot = context.buffer.isEmpty;
     context.buffer.write('SELECT ');
     var first = true;
     for (final MapEntry(key: expr, value: alias) in _columns.entries) {
@@ -202,7 +203,8 @@ final class SelectWithoutTables extends BaseSelectStatement<TypedResult>
       expr.writeInto(context);
       context.buffer.write(' ${context.identifier(alias)}');
     }
-    context.buffer.write(';');
+
+    if (isRoot) context.buffer.write(';');
   }
 
   GenerationContext _createContext() {

drift/lib/src/runtime/query_builder/statements/select/select_with_join.dart

@@ -7,7 +7,7 @@ part of '../../query_builder.dart';
 class JoinedSelectStatement<FirstT extends HasResultSet, FirstD>
     extends Query<FirstT, FirstD>
     with LimitContainerMixin, Selectable<TypedResult>
-    implements BaseSelectStatement {
+    implements BaseSelectStatement<TypedResult> {
   /// Used internally by drift, users should use [SimpleSelectStatement.join]
   /// instead.
   JoinedSelectStatement(super.database, super.table, this._joins,
@@ -36,6 +36,10 @@ class JoinedSelectStatement<FirstT extends HasResultSet, FirstD>
   /// here. They're just named in increasing order, so something like `AS c3`.
   final Map<Expression, String> _columnAliases = {};
 
+  /// Compound statements that have been added to this select statements, e.g.
+  /// through
+  final List<(_CompoundOperator, BaseSelectStatement)> _compounds = [];
+
   /// The tables this select statement reads from
   @visibleForOverriding
   @Deprecated('Use watchedTables on the generated context')
@@ -115,6 +119,149 @@ class JoinedSelectStatement<FirstT extends HasResultSet, FirstD>
     }
   }
 
+  void _addCompound(_CompoundOperator operator, BaseSelectStatement other) {
+    if (other is JoinedSelectStatement) {
+      if (other.limitExpr != null ||
+          other.orderByExpr != null ||
+          other._compounds.isNotEmpty) {
+        throw ArgumentError(
+            "Can't add compound query that has a limit or an order-by clause. "
+            'Also, the added query must hot have its own compound parts. Add  '
+            'the clauses and parts to the top-level parts instead.');
+      }
+    }
+
+    var columnsHere = _expandedColumns.iterator;
+    var otherColumns = other._expandedColumns.iterator;
+    var columnCount = 0;
+
+    while (columnsHere.moveNext()) {
+      if (!otherColumns.moveNext()) {
+        throw ArgumentError(
+            "Can't add select with fewer columns (added part has "
+            '$columnCount columns, the original source has more).');
+      }
+
+      var here = columnsHere.current;
+      var otherColumn = otherColumns.current;
+
+      if (here.$1.driftSqlType != otherColumn.$1.driftSqlType) {
+        throw ArgumentError(
+            "Can't add part because the column types at index $columnCount "
+            'differ.');
+      }
+
+      columnCount++;
+    }
+
+    if (otherColumns.moveNext()) {
+      throw ArgumentError(
+          "Can't add select with more columns (the original query has "
+          '$columnCount columns, the added part has more).');
+    }
+
+    _compounds.add((operator, other));
+  }
+
+  /// Appends the [other] statement as a `UNION` clause after this query.
+  ///
+  /// The database will run both queries and return all rows involved in either
+  /// query, removing duplicates. For this to work, this and [other] must have
+  /// compatible columns.
+  ///
+  /// The [other] query must not include a `LIMIT` or a `ORDER BY` clause.
+  /// Compound statements can only contain a single `LIMIT` and `ORDER BY`
+  /// clause at the end, which is set on the first statement (on which
+  /// [union] is called). Also, the [other] statement must not contain compound
+  /// parts on its own.
+  ///
+  /// As an example, consider a `todos` table of todo items referencing a
+  /// `categories` table used to group them. With that structure, it's possible
+  /// to compute the amount of todo items in each category, as well as the
+  /// amount of todo items not in a category in a single query:
+  ///
+  /// ```dart
+  ///   final count = subqueryExpression<int>(selectOnly(todos)
+  ///    ..addColumns([countAll()])
+  ///    ..where(todos.category.equalsExp(categories.id)));
+  ///  final countWithoutCategory = subqueryExpression<int>(db.selectOnly(todos)
+  ///        ..addColumns([countAll()])
+  ///        ..where(todos.category.isNull()));
+  ///
+  ///  final query = db.selectOnly(db.categories)
+  ///    ..addColumns([db.categories.description, count])
+  ///    ..groupBy([categories.id]);
+  ///  query.union(db.selectExpressions(
+  ///      [const Constant<String>(null), countWithoutCategory]));
+  /// ```
+  void union(BaseSelectStatement other) {
+    _addCompound(_CompoundOperator.union, other);
+  }
+
+  /// Appends the [other] statement as a `UNION ALL` clause after this query.
+  ///
+  /// The database will run both queries and return all rows involved in either
+  /// query. For this to work, this and [other] must have compatible columns.
+  ///
+  /// The [other] query must not include a `LIMIT` or a `ORDER BY` clause.
+  /// Compound statements can only contain a single `LIMIT` and `ORDER BY`
+  /// clause at the end, which is set on the first statement (on which
+  /// [unionAll] is called). Also, the [other] statement must not contain
+  /// compound parts on its own.
+  ///
+  /// As an example, consider a `todos` table of todo items referencing a
+  /// `categories` table used to group them. With that structure, it's possible
+  /// to compute the amount of todo items in each category, as well as the
+  /// amount of todo items not in a category in a single query:
+  ///
+  /// ```dart
+  ///   final count = subqueryExpression<int>(selectOnly(todos)
+  ///    ..addColumns([countAll()])
+  ///    ..where(todos.category.equalsExp(categories.id)));
+  ///  final countWithoutCategory = subqueryExpression<int>(db.selectOnly(todos)
+  ///        ..addColumns([countAll()])
+  ///        ..where(todos.category.isNull()));
+  ///
+  ///  final query = db.selectOnly(db.categories)
+  ///    ..addColumns([db.categories.description, count])
+  ///    ..groupBy([categories.id]);
+  ///  query.unionAll(db.selectExpressions(
+  ///      [const Constant<String>(null), countWithoutCategory]));
+  /// ```
+  void unionAll(BaseSelectStatement other) {
+    _addCompound(_CompoundOperator.unionAll, other);
+  }
+
+  /// Appends the [other] statement as a `EXCEPT` clause after this query.
+  ///
+  /// The database will run both queries and return all rows of the first query
+  /// that were not returned by [other]. For this to work, this and [other] must
+  /// have compatible columns.
+  ///
+  /// The [other] query must not include a `LIMIT` or a `ORDER BY` clause.
+  /// Compound statements can only contain a single `LIMIT` and `ORDER BY`
+  /// clause at the end, which is set on the first statement (on which
+  /// [except] is called). Also, the [other] statement must not contain
+  /// compound parts on its own.
+  void except(BaseSelectStatement other) {
+    _addCompound(_CompoundOperator.except, other);
+  }
+
+  /// Appends the [other] statement as a `INTERSECT` clause after this query.
+  ///
+  /// The database will run both queries and return all rows that were returned
+  /// by both queries. For this to work, this and [other] must have compatible
+  /// columns.
+  ///
+  /// The [other] query must not include a `LIMIT` or a `ORDER BY` clause.
+  /// Compound statements can only contain a single `LIMIT` and `ORDER BY`
+  /// clause at the end, which is set on the first statement (on which
+  /// [intersect] is called). Also, the [other] statement must not contain
+  /// compound parts on its own.
+  void intersect(BaseSelectStatement other) {
+    _addCompound(_CompoundOperator.intersect, other);
+  }
+
   @override
   void writeStartPart(GenerationContext ctx) {
     ctx.hasMultipleTables = true;
@@ -150,6 +297,28 @@ class JoinedSelectStatement<FirstT extends HasResultSet, FirstD>
     }
   }
 
+  @override
+  void writeInto(GenerationContext context) {
+    if (_compounds.isEmpty) {
+      super.writeInto(context);
+    } else {
+      // The order by and limit clauses must appear after the compounds.
+      super._writeInto(context, withOrderByAndLimit: false);
+
+      for (final (operator, statement) in _compounds) {
+        context.writeWhitespace();
+        context.buffer.write(operator.lexeme);
+        context.writeWhitespace();
+        statement.writeInto(context);
+      }
+
+      context.writeWhitespace();
+      orderByExpr?.writeInto(context);
+      context.writeWhitespace();
+      limitExpr?.writeInto(context);
+    }
+  }
+
   /// Applies the [predicate] as the where clause, which will be used to filter
   /// results.
   ///
@@ -401,3 +570,14 @@ extension JoinedSelectStatementAdditionalTables on JoinedSelectStatement {
               const []]) =>
       _watchWithAdditionalTables(tables);
 }
+
+enum _CompoundOperator {
+  union('UNION'),
+  unionAll('UNION ALL'),
+  intersect('INTERSECT'),
+  except('EXCEPT');
+
+  final String lexeme;
+
+  const _CompoundOperator(this.lexeme);
+}