Add createAggregateFunction method

Author: DallasHoff

docs/.vitepress/config.ts

@@ -69,6 +69,10 @@ export default defineConfig({
 						text: 'createScalarFunction',
 						link: '/api/createscalarfunction',
 					},
+					{
+						text: 'createAggregateFunction',
+						link: '/api/createaggregatefunction',
+					},
 					{
 						text: 'destroy',
 						link: '/api/destroy',

docs/api/createaggregatefunction.md

@@ -0,0 +1,52 @@
+# createAggregateFunction
+
+Create a SQL function that can be called from queries to combine multiple rows into a single result row.
+
+## Usage
+
+Access or destructure `createAggregateFunction` from the `SQLocal` client.
+
+```javascript
+import { SQLocal } from 'sqlocal';
+
+const { createAggregateFunction } = new SQLocal('database.sqlite3');
+```
+
+<!-- @include: ../.partials/initialization-note.md -->
+
+This method takes a string to name a custom SQL function as its first argument and an object containing two functions (`step` and `final`) as its second argument. After running `createAggregateFunction`, the aggregate function that you defined can be called from subsequent SQL queries. Arguments passed to the function in the SQL query will be passed to the JavaScript `step` function. The `step` function will run for every row in the SQL query. After each row is processed, the `final` function will run, and its return value will be passed back to SQLite to use to complete the query.
+
+This can be used to combine rows together in a query based on some custom logic. For example, the below aggregate function can be used to find the most common value for a column, such as the most common category used in a table of tasks.
+
+```javascript
+const values = new Map();
+
+await createAggregateFunction('mostCommon', {
+	step: (value) => {
+		const valueCount = values.get(value) ?? 0;
+		values.set(value, valueCount + 1);
+	},
+	final: () => {
+		const valueEntries = Array.from(values.entries());
+		const sortedEntries = valueEntries.sort((a, b) => b[1] - a[1]);
+		const mostCommonValue = sortedEntries[0][0];
+		values.clear();
+		return mostCommonValue;
+	},
+});
+
+await sql`SELECT mostCommon(category) AS mostCommonCategory FROM tasks`;
+```
+
+Aggregate functions can also be used in a query's HAVING clause to filter groups of rows. Here, we use the `mostCommon` function that we created in the previous example to find which days of the week have "Cleaning" as the most common category of task.
+
+```javascript
+await sql`
+  SELECT dayOfWeek
+  FROM tasks
+  GROUP BY dayOfWeek
+  HAVING mostCommon(category) = 'Cleaning'
+`;
+```
+
+<!-- @include: ../.partials/functions-note.md -->

src/client.ts

@@ -11,6 +11,7 @@ import type {
 	StatementInput,
 	Transaction,
 	DatabasePath,
+	AggregateUserFunction,
 } from './types.js';
 import type {
 	BatchMessage,
@@ -391,9 +392,12 @@ export class SQLocal {
 		func: ScalarUserFunction['func']
 	): Promise<void> => {
 		const key = `_sqlocal_func_${funcName}`;
+		const attachFunction = () => {
+			this.proxy[key] = func;
+		};
 
 		if (this.proxy === globalThis) {
-			this.proxy[key] = func;
+			attachFunction();
 		}
 
 		await this.createQuery({
@@ -403,7 +407,32 @@ export class SQLocal {
 		});
 
 		if (this.proxy !== globalThis) {
-			this.proxy[key] = func;
+			attachFunction();
+		}
+	};
+
+	createAggregateFunction = async (
+		funcName: string,
+		func: AggregateUserFunction['func']
+	): Promise<void> => {
+		const key = `_sqlocal_func_${funcName}`;
+		const attachFunction = () => {
+			this.proxy[`${key}_step`] = func.step;
+			this.proxy[`${key}_final`] = func.final;
+		};
+
+		if (this.proxy === globalThis) {
+			attachFunction();
+		}
+
+		await this.createQuery({
+			type: 'function',
+			functionName: funcName,
+			functionType: 'aggregate',
+		});
+
+		if (this.proxy !== globalThis) {
+			attachFunction();
 		}
 	};
 

src/drivers/sqlite-memory-driver.ts

@@ -87,11 +87,24 @@ export class SQLiteMemoryDriver implements SQLocalDriver {
 	async createFunction(fn: UserFunction): Promise<void> {
 		if (!this.db) throw new Error('Driver not initialized');
 
-		this.db.createFunction({
-			name: fn.name,
-			xFunc: (_: number, ...args: any[]) => fn.func(...args),
-			arity: -1,
-		});
+		switch (fn.type) {
+			case 'callback':
+			case 'scalar':
+				this.db.createFunction({
+					name: fn.name,
+					xFunc: (_: number, ...args: any[]) => fn.func(...args),
+					arity: -1,
+				});
+				break;
+			case 'aggregate':
+				this.db.createFunction({
+					name: fn.name,
+					xStep: (_: number, ...args: any[]) => fn.func.step(...args),
+					xFinal: (_: number, ...args: any[]) => fn.func.final(...args),
+					arity: -1,
+				});
+				break;
+		}
 	}
 
 	async import(

src/processor.ts

@@ -266,38 +266,51 @@ export class SQLocalProcessor {
 	protected createUserFunction = async (
 		message: FunctionMessage
 	): Promise<void> => {
-		const { functionName, functionType, queryKey } = message;
-		let func;
+		const { functionName: name, functionType: type, queryKey } = message;
+		let fn: UserFunction;
 
-		if (this.userFunctions.has(functionName)) {
+		if (this.userFunctions.has(name)) {
 			this.emitMessage({
 				type: 'error',
 				error: new Error(
-					`A user-defined function with the name "${functionName}" has already been created for this SQLocal instance.`
+					`A user-defined function with the name "${name}" has already been created for this SQLocal instance.`
 				),
 				queryKey,
 			});
 			return;
 		}
 
-		if (functionType === 'callback') {
-			func = (...args: any[]) => {
-				this.emitMessage({
-					type: 'callback',
-					name: functionName,
-					args: args,
-				});
-			};
-		} else {
-			func = this.proxy[`_sqlocal_func_${functionName}`];
+		switch (type) {
+			case 'callback':
+				fn = {
+					type,
+					name,
+					func: (...args: any[]) => {
+						this.emitMessage({ type: 'callback', name, args });
+					},
+				};
+				break;
+			case 'scalar':
+				fn = {
+					type,
+					name,
+					func: this.proxy[`_sqlocal_func_${name}`],
+				};
+				break;
+			case 'aggregate':
+				fn = {
+					type,
+					name,
+					func: {
+						step: this.proxy[`_sqlocal_func_${name}_step`],
+						final: this.proxy[`_sqlocal_func_${name}_final`],
+					},
+				};
+				break;
 		}
 
 		try {
-			await this.initUserFunction({
-				type: functionType,
-				name: functionName,
-				func,
-			});
+			await this.initUserFunction(fn);
 			this.emitMessage({
 				type: 'success',
 				queryKey,

src/types.ts

@@ -123,7 +123,10 @@ export type DatabaseInfo = {
 
 // User functions
 
-export type UserFunction = CallbackUserFunction | ScalarUserFunction;
+export type UserFunction =
+	| CallbackUserFunction
+	| ScalarUserFunction
+	| AggregateUserFunction;
 export type CallbackUserFunction = {
 	type: 'callback';
 	name: string;
@@ -134,3 +137,11 @@ export type ScalarUserFunction = {
 	name: string;
 	func: (...args: any[]) => any;
 };
+export type AggregateUserFunction = {
+	type: 'aggregate';
+	name: string;
+	func: {
+		step: (...args: any[]) => void;
+		final: (...args: any[]) => any;
+	};
+};

test/create-aggregate-function.test.ts

@@ -0,0 +1,69 @@
+import { afterEach, beforeAll, beforeEach, describe, expect, it } from 'vitest';
+import { SQLocal } from '../src/index.js';
+
+describe.each([
+	{ type: 'opfs', path: 'create-aggregate-function-test.sqlite3' },
+	{ type: 'memory', path: ':memory:' },
+	{ type: 'local', path: ':localStorage:' },
+	{ type: 'session', path: ':sessionStorage:' },
+])('createAggregateFunction ($type)', ({ path }) => {
+	const { sql, createAggregateFunction } = new SQLocal(path);
+
+	beforeAll(async () => {
+		const values = new Map<unknown, number>();
+
+		await createAggregateFunction('mostCommon', {
+			step: (value: unknown) => {
+				const valueCount = values.get(value) ?? 0;
+				values.set(value, valueCount + 1);
+			},
+			final: () => {
+				const valueEntries = Array.from(values.entries());
+				const sortedEntries = valueEntries.sort((a, b) => b[1] - a[1]);
+				const mostCommonValue = sortedEntries[0][0];
+				values.clear();
+				return mostCommonValue;
+			},
+		});
+	});
+
+	beforeEach(async () => {
+		await sql`CREATE TABLE nums (num REAL NOT NULL)`;
+	});
+
+	afterEach(async () => {
+		await sql`DROP TABLE nums`;
+	});
+
+	it('should create and use aggregate function in SELECT clause', async () => {
+		await sql`INSERT INTO nums (num) VALUES (0), (3), (2), (7), (3), (1), (5), (3), (3), (2)`;
+
+		const results = await sql`SELECT mostCommon(num) AS mostCommon FROM nums`;
+
+		expect(results).toEqual([{ mostCommon: 3 }]);
+	});
+
+	it('should create and use aggregate function in HAVING clause', async () => {
+		await sql`INSERT INTO nums (num) VALUES (1), (2), (2), (2), (4), (5), (5), (6)`;
+
+		const results = await sql`
+			SELECT mod(num, 2) AS isOdd
+			FROM nums
+			GROUP BY isOdd
+			HAVING mostCommon(num) = 5
+		`;
+
+		expect(results).toEqual([{ isOdd: 1 }]);
+	});
+
+	it('should not replace an existing implementation', async () => {
+		const createBadFn = async () => {
+			await createAggregateFunction('mostCommon', {
+				step: () => {},
+				final: () => 0,
+			});
+		};
+
+		await expect(createBadFn).rejects.toThrowError();
+	});
+});

test/create-scalar-function.test.ts

@@ -17,7 +17,7 @@ describe.each([
 		await sql`DROP TABLE nums`;
 	});
 
-	it('should create and use scalar function in columns clause', async () => {
+	it('should create and use scalar function in SELECT clause', async () => {
 		await createScalarFunction('double', (num: number) => num * 2);
 
 		const createBadFn = async () =>
@@ -26,7 +26,7 @@ describe.each([
 
 		await sql`INSERT INTO nums (num) VALUES (0), (2), (3.5), (-11.11)`;
 
-		const results = await sql`SELECT num, double(num) as doubled FROM nums`;
+		const results = await sql`SELECT num, double(num) AS doubled FROM nums`;
 
 		expect(results).toEqual([
 			{ num: 0, doubled: 0 },
@@ -36,7 +36,7 @@ describe.each([
 		]);
 	});
 
-	it('should create and use scalar function in where clause', async () => {
+	it('should create and use scalar function in WHERE clause', async () => {
 		await createScalarFunction('isEven', (num: number) => num % 2 === 0);
 
 		await sql`INSERT INTO nums (num) VALUES (2), (3), (4), (5), (6)`;
@@ -74,8 +74,8 @@ describe.each([
 		await db1.createScalarFunction('addTax', (num: number) => num * 1.06);
 		await db2.createScalarFunction('addTax', (num: number) => num * 1.07);
 
-		const [result1] = await db1.sql`SELECT addTax(2) as withTax`;
-		const [result2] = await db2.sql`SELECT addTax(2) as withTax`;
+		const [result1] = await db1.sql`SELECT addTax(2) AS withTax`;
+		const [result2] = await db2.sql`SELECT addTax(2) AS withTax`;
 
 		expect(result1.withTax).toBe(2.12);
 		expect(result2.withTax).toBe(2.14);