perf(cassandra): memoize CQL statement metadata per query (#2155)

fluxer_api/src/api/database/CassandraDevLogger.ts

@@ -2,6 +2,7 @@
 
 import {Logger} from '../Logger';
 import {getIsDev} from './CassandraMetaRegistry';
+import {getStatementMeta} from './CassandraTypes';
 
 const colors = {
 	reset: '\x1b[0m',
@@ -67,16 +68,6 @@ function formatParams(params: Record<string, unknown>): string {
 	return entries.map(([k, v]) => `  ${colors.blue}:${k}${colors.reset} = ${formatValue(v)}`).join('\n');
 }
 
-export function getQueryType(cql: string): string {
-	const trimmed = cql.trim().toUpperCase();
-	if (trimmed.startsWith('SELECT')) return 'SELECT';
-	if (trimmed.startsWith('INSERT')) return 'INSERT';
-	if (trimmed.startsWith('UPDATE')) return 'UPDATE';
-	if (trimmed.startsWith('DELETE')) return 'DELETE';
-	if (trimmed.startsWith('BEGIN BATCH')) return 'BATCH';
-	return 'QUERY';
-}
-
 function formatCql(cql: string): string {
 	return cql
 		.replace(/\s+/g, ' ')
@@ -128,7 +119,7 @@ export function logBatch(
 	];
 	for (let i = 0; i < queries.length; i++) {
 		const {query, params} = queries[i];
-		const queryType = getQueryType(query);
+		const queryType = getStatementMeta(query).type;
 		const typeColors: Record<string, string> = {
 			SELECT: colors.cyan,
 			INSERT: colors.green,

fluxer_api/src/api/database/CassandraQueryExecution.ts

@@ -4,13 +4,13 @@ import {ServiceUnavailableError} from '@fluxer/errors/src/domains/core/ServiceUn
 import {getClient} from '@pkgs/cassandra/src/Client';
 import cassandra from 'cassandra-driver';
 import {Logger} from '../Logger';
-import {getQueryType, logBatch, logQuery} from './CassandraDevLogger';
+import {logBatch, logQuery} from './CassandraDevLogger';
 import {getIsDev} from './CassandraMetaRegistry';
 import type {CassandraParams, KvQueryMeta, PreparedQuery, QueryTemplate} from './CassandraTypes';
 import {
 	assertNoUndefinedParams,
 	chunkArray,
-	isUnsafePreparedStatement,
+	getStatementMeta,
 	normalizeExecuteArgs,
 	normalizeInParams,
 } from './CassandraTypes';
@@ -98,8 +98,9 @@ export async function executeQuery<T = Record<string, unknown>, P extends Cassan
 	params?: P,
 ): Promise<Array<T>> {
 	const {cql, params: boundRaw} = normalizeExecuteArgs(queryOrPrepared, params);
-	const bound = normalizeInParams(cql, boundRaw);
-	if (isUnsafePreparedStatement(cql)) {
+	const meta = getStatementMeta(cql);
+	const bound = normalizeInParams(meta, boundRaw);
+	if (meta.unsafe) {
 		throw new Error('Cannot prepare a statement that looks like `SELECT *`');
 	}
 	assertNoUndefinedParams(bound as Record<string, unknown>);
@@ -111,14 +112,14 @@ export async function executeQuery<T = Record<string, unknown>, P extends Cassan
 			kvMeta: typeof queryOrPrepared === 'string' ? undefined : queryOrPrepared.kvMeta,
 		});
 	}
-	const startTime = getIsDev() ? performance.now() : Date.now();
-	const queryType = getQueryType(cql);
+	const isDev = getIsDev();
+	const startTime = isDev ? performance.now() : 0;
 	try {
 		const result = await getClient().execute(cql, bound, {prepare: true});
-		const rows = await collectSelectRows<T>(queryType, result);
-		if (getIsDev()) {
+		const rows = await collectSelectRows<T>(meta.type, result);
+		if (isDev) {
 			const durationMs = performance.now() - startTime;
-			logQuery(queryType, cql, bound as Record<string, unknown>, durationMs, rows.length);
+			logQuery(meta.type, cql, bound as Record<string, unknown>, durationMs, rows.length);
 		}
 		return rows;
 	} catch (err: unknown) {
@@ -166,8 +167,9 @@ export async function fetchPage<T = Record<string, unknown>, P extends Cassandra
 	},
 ): Promise<PagedQueryResult<T>> {
 	const {cql, params: boundRaw} = normalizeExecuteArgs(queryOrPrepared, params);
-	const bound = normalizeInParams(cql, boundRaw);
-	if (isUnsafePreparedStatement(cql)) {
+	const meta = getStatementMeta(cql);
+	const bound = normalizeInParams(meta, boundRaw);
+	if (meta.unsafe) {
 		throw new Error('Cannot prepare a statement that looks like `SELECT *`');
 	}
 	assertNoUndefinedParams(bound as Record<string, unknown>);
@@ -248,7 +250,7 @@ interface BatchQuery {
 async function executeBatch(queries: Array<BatchQuery>, atomic = true): Promise<void> {
 	if (queries.length === 0) return;
 	for (const {query} of queries) {
-		if (isUnsafePreparedStatement(query)) {
+		if (getStatementMeta(query).unsafe) {
 			throw new Error('Cannot prepare a statement that looks like `SELECT *`');
 		}
 	}
@@ -265,16 +267,20 @@ async function executeBatch(queries: Array<BatchQuery>, atomic = true): Promise<
 		logged: atomic,
 		counter: false,
 	};
-	const startTime = getIsDev() ? performance.now() : 0;
+	const isDev = getIsDev();
+	const startTime = isDev ? performance.now() : 0;
 	try {
 		await getClient().batch(
-			queries.map(({query, params}) => ({query, params: normalizeInParams(query, params as CassandraParams)})),
+			queries.map(({query, params}) => ({
+				query,
+				params: normalizeInParams(getStatementMeta(query), params as CassandraParams),
+			})),
 			options,
 		);
 	} catch (err: unknown) {
 		throw mapCassandraDriverError(err);
 	}
-	if (getIsDev()) {
+	if (isDev) {
 		const durationMs = performance.now() - startTime;
 		logBatch(queries, durationMs);
 	}

fluxer_api/src/api/database/CassandraTypes.test.ts

@@ -0,0 +1,74 @@
+// SPDX-License-Identifier: AGPL-3.0-or-later
+
+import {describe, expect, it} from 'vitest';
+import {BatchBuilder, executeQuery, fetchPage} from './CassandraQueryExecution';
+import {getStatementMeta, normalizeInParams} from './CassandraTypes';
+
+describe('getStatementMeta', () => {
+	it('flags SELECT * statements as unsafe', () => {
+		expect(getStatementMeta('SELECT * FROM messages WHERE channel_id = :channel_id').unsafe).toBe(true);
+		expect(getStatementMeta('  select   *   FROM messages').unsafe).toBe(true);
+		expect(getStatementMeta('SELECT\n\t*\nFROM messages').unsafe).toBe(true);
+	});
+
+	it('does not flag projected selects or other statements as unsafe', () => {
+		expect(getStatementMeta('SELECT id, content FROM messages WHERE id = :id').unsafe).toBe(false);
+		expect(getStatementMeta('SELECT count(*) FROM messages').unsafe).toBe(false);
+		expect(getStatementMeta('UPDATE messages SET content = :content WHERE id = :id').unsafe).toBe(false);
+		expect(getStatementMeta('SELECT').unsafe).toBe(false);
+		expect(getStatementMeta('').unsafe).toBe(false);
+	});
+
+	it('derives the statement type', () => {
+		expect(getStatementMeta('  SELECT id FROM messages').type).toBe('SELECT');
+		expect(getStatementMeta('insert into messages (id) VALUES (:id)').type).toBe('INSERT');
+		expect(getStatementMeta('UPDATE messages SET content = :content').type).toBe('UPDATE');
+		expect(getStatementMeta('DELETE FROM messages WHERE id = :id').type).toBe('DELETE');
+		expect(getStatementMeta('BEGIN BATCH APPLY BATCH').type).toBe('BATCH');
+		expect(getStatementMeta('TRUNCATE messages').type).toBe('QUERY');
+	});
+
+	it('collects distinct IN parameter names in order', () => {
+		expect(getStatementMeta('SELECT id FROM messages WHERE id IN :ids AND author_id IN (:authors)').inParams).toEqual([
+			'ids',
+			'authors',
+		]);
+		expect(getStatementMeta('SELECT id FROM messages WHERE id IN :ids OR other_id IN :ids').inParams).toEqual(['ids']);
+		expect(getStatementMeta('SELECT id FROM messages WHERE id = :id').inParams).toEqual([]);
+	});
+
+	it('reuses the memoized entry for the same statement', () => {
+		const cql = 'SELECT id FROM memoized_rows WHERE id IN :ids';
+		const first = getStatementMeta(cql);
+		expect(getStatementMeta(cql)).toBe(first);
+		expect(getStatementMeta(`${cql} `)).not.toBe(first);
+	});
+});
+
+describe('normalizeInParams', () => {
+	it('expands Set values bound to IN parameters into arrays', () => {
+		const meta = getStatementMeta('SELECT id FROM messages WHERE id IN :ids');
+		const params = {ids: new Set(['a', 'b'])};
+		const normalized = normalizeInParams(meta, params);
+		expect(normalized).not.toBe(params);
+		expect(normalized.ids).toEqual(['a', 'b']);
+	});
+
+	it('returns the original params when there is nothing to expand', () => {
+		const withoutIn = getStatementMeta('SELECT id FROM messages WHERE id = :id');
+		const params = {id: 'a'};
+		expect(normalizeInParams(withoutIn, params)).toBe(params);
+		const withIn = getStatementMeta('SELECT id FROM messages WHERE id IN :ids');
+		const arrayParams = {ids: ['a', 'b']};
+		expect(normalizeInParams(withIn, arrayParams)).toBe(arrayParams);
+	});
+});
+
+describe('unsafe statement guard', () => {
+	it('rejects SELECT * through every execution entrypoint', async () => {
+		const message = 'Cannot prepare a statement that looks like `SELECT *`';
+		await expect(executeQuery('SELECT * FROM messages', {})).rejects.toThrow(message);
+		await expect(fetchPage('SELECT * FROM messages', {}, {pageSize: 1})).rejects.toThrow(message);
+		await expect(new BatchBuilder().add('SELECT * FROM messages', {}).execute()).rejects.toThrow(message);
+	});
+});

fluxer_api/src/api/database/CassandraTypes.ts

@@ -215,11 +215,15 @@ export function normalizeExecuteArgs<P extends CassandraParams>(
 	return queryOrPrepared;
 }
 
-const IN_PARAM_CACHE = new Map<string, Array<string>>();
+interface CqlStatementMeta {
+	unsafe: boolean;
+	type: string;
+	inParams: ReadonlyArray<string>;
+}
 
-function getInParamNames(cql: string): Array<string> {
-	const cached = IN_PARAM_CACHE.get(cql);
-	if (cached) return cached;
+const STATEMENT_META_CACHE = new Map<string, CqlStatementMeta>();
+
+function computeInParamNames(cql: string): Array<string> {
 	const regex = /\bIN\s*\(?\s*:([A-Za-z_][A-Za-z0-9_]*)/g;
 	const names: Array<string> = [];
 	const seen = new Set<string>();
@@ -230,16 +234,42 @@ function getInParamNames(cql: string): Array<string> {
 		seen.add(name);
 		names.push(name);
 	}
-	IN_PARAM_CACHE.set(cql, names);
 	return names;
 }
 
-export function normalizeInParams<P extends CassandraParams>(cql: string, params: P): P {
-	const inParams = getInParamNames(cql);
-	if (inParams.length === 0) return params;
+function computeQueryType(trimmed: string): string {
+	const upper = trimmed.toUpperCase();
+	if (upper.startsWith('SELECT')) return 'SELECT';
+	if (upper.startsWith('INSERT')) return 'INSERT';
+	if (upper.startsWith('UPDATE')) return 'UPDATE';
+	if (upper.startsWith('DELETE')) return 'DELETE';
+	if (upper.startsWith('BEGIN BATCH')) return 'BATCH';
+	return 'QUERY';
+}
+
+function computeStatementMeta(cql: string): CqlStatementMeta {
+	const trimmed = cql.trim();
+	const tokens = trimmed.split(/\s+/);
+	return {
+		unsafe: tokens.length >= 2 && tokens[0].toLowerCase() === 'select' && tokens[1] === '*',
+		type: computeQueryType(trimmed),
+		inParams: computeInParamNames(cql),
+	};
+}
+
+export function getStatementMeta(cql: string): CqlStatementMeta {
+	const cached = STATEMENT_META_CACHE.get(cql);
+	if (cached) return cached;
+	const meta = computeStatementMeta(cql);
+	STATEMENT_META_CACHE.set(cql, meta);
+	return meta;
+}
+
+export function normalizeInParams<P extends CassandraParams>(meta: CqlStatementMeta, params: P): P {
+	if (meta.inParams.length === 0) return params;
 	let changed = false;
 	const out: Record<string, CassandraParam> = {...params};
-	for (const paramName of inParams) {
+	for (const paramName of meta.inParams) {
 		const value = out[paramName];
 		if (value instanceof Set) {
 			out[paramName] = Array.from(value.values());
@@ -249,11 +279,6 @@ export function normalizeInParams<P extends CassandraParams>(cql: string, params
 	return (changed ? (out as P) : params) as P;
 }
 
-export function isUnsafePreparedStatement(query: string): boolean {
-	const tokens = query.trim().split(/\s+/);
-	return tokens.length >= 2 && tokens[0].toLowerCase() === 'select' && tokens[1] === '*';
-}
-
 function hasUndefinedDeep(value: unknown): boolean {
 	if (value === undefined) return true;
 	if (value === null) return false;