feat: adds support for .count() in the query builder

packages/backend/src/execution/executors/PgExecutor/PgExecutor.test.ts

@@ -48,6 +48,29 @@ describe('PgExecutor', () => {
     `);
     });
 
+    it('.count() ', async () => {
+        const q = from('film')
+            .where({
+                film_id: { in: [1, 2, 3] },
+            })
+            .groupBy()
+            .count();
+        const { sql } = executor.compile(q);
+
+        expect(sql).toMatchInlineSnapshot(`
+      "select
+        count(1) as "count"
+      from
+        "public"."film" "public::film"
+      where
+        ("public::film".film_id = any ($1))"
+    `);
+
+        const result = await executor.execute(q, executeProps);
+
+        expect(result).toEqual([{ count: 3 }]);
+    });
+
     it('Film table SynthQL query executes to expected result', async () => {
         const result = await executor.execute(q1, executeProps);
 

packages/backend/src/execution/executors/PgExecutor/PgExecutor.ts

@@ -1,8 +1,6 @@
 import { Pool, PoolClient } from 'pg';
 import { format } from 'sql-formatter';
 import { splitQueryAtBoundary } from '../../../query/splitQueryAtBoundary';
-import { ColumnRef } from '../../../refs/ColumnRef';
-import { RefContext, createRefContext } from '../../../refs/RefContext';
 import { AnyQuery } from '../../../types';
 import { QueryExecutor } from '../../types';
 import { QueryProviderExecutor } from '../QueryProviderExecutor';

packages/backend/src/execution/executors/PgExecutor/queryBuilder/AggregationsSelection.ts

@@ -0,0 +1,40 @@
+import { AnyQuery } from '../../../../types';
+import { Selection } from './types';
+import { SqlBuilder } from './exp';
+
+// TODO(fhur): this is actually only supporting count(*) for now.
+export class AggregationsSelection implements Selection {
+    public constructor() {}
+
+    static fromQuery(
+        rootQuery: AnyQuery,
+        defaultSchema: string,
+    ): AggregationsSelection[] {
+        if (rootQuery.aggregates === undefined) {
+            return [];
+        }
+        if (Object.keys(rootQuery.aggregates).length === 0) {
+            return [];
+        }
+        return [new AggregationsSelection()];
+    }
+
+    toSql() {
+        return new SqlBuilder().addAs(['as', ['fn', 'count', '1'], 'count']);
+    }
+
+    extractFromRow(row: any, target: any): void {
+        // TODO assert prescence
+        const tmp = row['count'];
+        if (typeof tmp !== 'string') {
+            throw new Error(
+                `Expected count to be a string, got: ${JSON.stringify(tmp)}`,
+            );
+        }
+
+        // Note technically correct, but we're not going to support counting bigints
+        // this is only an issue if you return more than Integer.MAX_VALUE rows,
+        // which is unlikely.
+        target['count'] = parseInt(tmp);
+    }
+}

packages/backend/src/execution/executors/PgExecutor/queryBuilder/createAugmentedQuery.ts

@@ -6,6 +6,7 @@ import { SelectionJsonbAgg } from './SelectionJsonbAgg';
 import { TableRef } from '../../../../refs/TableRef';
 import { ColumnRef } from '../../../../refs/ColumnRef';
 import { Join, Selection } from './types';
+import { AggregationsSelection } from './AggregationsSelection';
 
 export interface AugmentedQuery {
     selection: Selection[];
@@ -29,13 +30,17 @@ export function createAugmentedQuery(
         rootQuery,
         defaultSchema,
     );
+    const aggregates = AggregationsSelection.fromQuery(
+        rootQuery,
+        defaultSchema,
+    );
     const wheres = collectWhere(rootQuery, defaultSchema);
     const joins = collectJoins(rootQuery, defaultSchema);
     const groupingColumns =
         rootQuery.groupBy?.map((col) => rootTable.column(col)) ?? [];
 
     return {
-        selection: selectionColumn.concat(jsonbAggColumn),
+        selection: selectionColumn.concat(jsonbAggColumn).concat(aggregates),
         rootQuery,
         rootTable,
         joins,

packages/backend/src/execution/executors/PgExecutor/queryBuilder/exp.ts

@@ -229,7 +229,6 @@ export class SqlBuilder {
     }
 
     addOperator(op: BinaryOperator) {
-        const unknownOp = op as unknown;
         if (!OPERATORS.includes(op as any)) {
             throw new Error(`Invalid operator: ${op}`);
         }

packages/backend/src/tests/e2e/nx1.test.ts

@@ -1,4 +1,4 @@
-import { Where, col } from '@synthql/queries';
+import { QueryResult, Where, col } from '@synthql/queries';
 import { describe, expect, test } from 'vitest';
 import { collectLast } from '../..';
 import { execute } from '../../execution/execute';
@@ -166,4 +166,10 @@ describe('n x 1', () => {
             expect(c.expected).toEqual(c.expected);
         });
     });
+
+    test('x', async () => {
+        const q = from('film').count();
+
+        type x = QueryResult<DB, typeof q>;
+    });
 });

packages/backend/src/tests/propertyBased/arbitraries/ArbitraryQueryBuilder.ts

@@ -19,6 +19,8 @@ export class ArbitraryQueryBuilder<DB> {
         private cardinalities: Cardinality[] = ['many', 'maybe', 'one'],
         private tables: Table<DB>[] = getTableNames(schema) as Table<DB>[],
         private hasResults: boolean = true,
+        private groupBy?: string[],
+        private skipLimit?: boolean,
     ) {}
 
     /**
@@ -34,6 +36,19 @@ export class ArbitraryQueryBuilder<DB> {
             cardinality,
             this.tables,
             this.hasResults,
+            this.groupBy,
+            this.skipLimit,
+        );
+    }
+
+    withGroupBy(...groupBy: string[]) {
+        return new ArbitraryQueryBuilder<DB>(
+            this.schema,
+            this.cardinalities,
+            this.tables,
+            this.hasResults,
+            groupBy,
+            this.skipLimit,
         );
     }
 
@@ -43,6 +58,8 @@ export class ArbitraryQueryBuilder<DB> {
             this.cardinalities,
             tables,
             this.hasResults,
+            this.groupBy,
+            this.skipLimit,
         );
     }
 
@@ -57,6 +74,8 @@ export class ArbitraryQueryBuilder<DB> {
             this.cardinalities,
             this.tables,
             false,
+            this.groupBy,
+            this.skipLimit,
         );
     }
 
@@ -71,6 +90,19 @@ export class ArbitraryQueryBuilder<DB> {
             this.cardinalities,
             this.tables,
             true,
+            this.groupBy,
+            this.skipLimit,
+        );
+    }
+
+    withNoLimit() {
+        return new ArbitraryQueryBuilder<DB>(
+            this.schema,
+            this.cardinalities,
+            this.tables,
+            this.hasResults,
+            this.groupBy,
+            true,
         );
     }
 
@@ -115,14 +147,20 @@ export class ArbitraryQueryBuilder<DB> {
     }
 
     private arbGroupBy(tableName: Table<DB>): fc.Arbitrary<string[]> {
-        return fc.constant(getPrimaryKeyColumns(this.schema, tableName));
+        if (this.groupBy === undefined) {
+            return fc.constant(getPrimaryKeyColumns(this.schema, tableName));
+        }
+        return fc.constant(this.groupBy);
     }
 
     private arbCardinality(): fc.Arbitrary<Cardinality> {
         return fc.constantFrom(...this.cardinalities);
     }
 
-    private arbLimit(): fc.Arbitrary<number> {
+    private arbLimit(): fc.Arbitrary<number | undefined> {
+        if (this.skipLimit) {
+            return fc.constant(undefined);
+        }
         if (!this.hasResults) {
             return fc.constant(0);
         }

packages/backend/src/tests/propertyBased/properties/count.test.ts

@@ -0,0 +1,60 @@
+import { test } from '@fast-check/vitest';
+import { ArbitraryQueryBuilder } from '../arbitraries/ArbitraryQueryBuilder';
+import { queryEngine } from '../../queryEngine';
+import { expect } from 'vitest';
+import { Query } from '@synthql/queries';
+import { DB } from '../../generated';
+
+const queryBuilder = ArbitraryQueryBuilder.fromPagila();
+
+const numRuns = 100;
+const timeout = numRuns * 1000;
+const endOnFailure = true;
+
+test.prop(
+    [
+        queryBuilder
+            .withCardinality('many')
+            .withNoLimit()
+            .withSomeResults()
+            .build(),
+    ],
+    {
+        verbose: true,
+        numRuns,
+        timeout,
+        endOnFailure,
+        seed: 1308343585,
+        path: '2',
+    },
+)(
+    'execute(query).length should equal execute(query.count()).count',
+    async (query) => {
+        const queryResult = (await queryEngine.executeAndWait(
+            query,
+        )) as Array<any>;
+
+        const countQuery: Query<DB> = {
+            ...query,
+            select: {},
+            aggregates: {
+                count: {
+                    type: 'fn',
+                    fn: 'count',
+                    args: [1],
+                },
+            },
+            groupBy: [],
+            cardinality: 'one',
+        };
+
+        const countResult = (await queryEngine.executeAndWait(
+            countQuery,
+        )) as any;
+
+        expect(typeof countResult.count).toEqual('number');
+
+        expect(queryResult.length).toEqual(countResult.count);
+    },
+    timeout,
+);