fix: update custom tables replacement function

Author: yannamsellem

src/lib/olap-engine/index.ts

@@ -43,4 +43,4 @@ export const engine: OLAPEngine =
 		? new RemoteEngine()
 		: new LocalEngine();
 
-export { applySlugs } from './tables';
+export { applyCustomTable } from './tables';

src/lib/olap-engine/tables.ts

@@ -3,17 +3,60 @@ interface Source {
 	path: string;
 }
 
-export function applySlugs(query: string, sources: Source[]) {
-	for (const source of sources) {
-		query = query.replace(new RegExp(`(from|FROM)[ \n\t]+(${source.slug})`, 'g'), (match) =>
-			match.replace(source.slug, `${source.path} ${source.slug}`)
+const SQL_KEYWORDS = new Set([
+	'WHERE',
+	'ON',
+	'GROUP',
+	'ORDER',
+	'LIMIT',
+	'OFFSET',
+	'HAVING',
+	'SET',
+	'LEFT',
+	'RIGHT',
+	'INNER',
+	'OUTER',
+	'CROSS',
+	'JOIN',
+	'UNION',
+	'RETURNING',
+	'VALUES',
+	';'
+]);
+
+export function applyCustomTable(query: string, sources: Source[]) {
+	let newQuery = query;
+
+	for (const { slug, path } of sources) {
+		const regex = new RegExp(
+			`(\\bFROM\\b|\\bJOIN\\b|\\bUPDATE\\b|\\bINTO\\b)\\s+(?:(\\w+)\\.)?(${slug})\\b(\\s+(?:AS\\s+)?\\w+)?`,
+			'gi'
+		);
+
+		newQuery = newQuery.replace(
+			regex,
+			(_fullMatch, keyword, _schema, foundTable, potentialAliasChunk) => {
+				let finalAliasPart = '';
+
+				if (potentialAliasChunk) {
+					const rawAlias = potentialAliasChunk.trim().replace(/^AS\s+/i, '');
+					if (SQL_KEYWORDS.has(rawAlias.toUpperCase())) {
+						finalAliasPart = ` ${foundTable} ${potentialAliasChunk}`;
+					} else {
+						finalAliasPart = potentialAliasChunk;
+					}
+				} else {
+					finalAliasPart = ` ${foundTable}`;
+				}
+				return `${keyword} ${path}${finalAliasPart}`;
+			}
 		);
 
-		query = query.replace(
-			new RegExp(`(describe|DESCRIBE)([ \n\t]+(table|TABLE))?[ \n\t]+(${source.slug})`, 'g'),
-			(match) => match.replace(source.slug, `${source.path}`)
+		newQuery = newQuery.replace(
+			new RegExp(`(DESCRIBE)\\s+(?:TABLE\\s+)?(${slug})\\b`, 'gi'),
+			`$1 ${path}`
 		);
 	}
 
-	return query;
+	return newQuery;
 }

src/routes/+page.svelte

@@ -33,7 +33,7 @@
 	import Sparkles from '$lib/icons/Sparkles.svelte';
 	import Stop from '$lib/icons/Stop.svelte';
 	import type { Table } from '$lib/olap-engine';
-	import { applySlugs, engine, type OLAPResponse } from '$lib/olap-engine';
+	import { applyCustomTable, engine, type OLAPResponse } from '$lib/olap-engine';
 	import { SQLiteChatsRepository, type ChatsRepository } from '$lib/repositories/chats';
 	import {
 		SQLiteHistoryRepository,
@@ -100,7 +100,9 @@
 
 			cached = false;
 			abortController = new AbortController();
-			response = await engine.exec(applySlugs(query, sources), { signal: abortController.signal });
+			response = await engine.exec(applyCustomTable(query, sources), {
+				signal: abortController.signal
+			});
 			await cache.set(query, response);
 		} finally {
 			loading = false;