diff --git a/internal-packages/tsql/package.json b/internal-packages/tsql/package.json index 5ebf1f130..7baf40726 100644 --- a/internal-packages/tsql/package.json +++ b/internal-packages/tsql/package.json @@ -6,8 +6,10 @@ "types": "./src/index.ts", "type": "module", "dependencies": { + "@internal/clickhouse": "workspace:*", "@trigger.dev/core": "workspace:*", - "antlr4ts": "0.5.0-alpha.4" + "antlr4ts": "0.5.0-alpha.4", + "zod": "3.25.76" }, "scripts": { "typecheck": "tsc --noEmit", diff --git a/internal-packages/tsql/src/index.ts b/internal-packages/tsql/src/index.ts index 6580918fb..bcb0476b3 100644 --- a/internal-packages/tsql/src/index.ts +++ b/internal-packages/tsql/src/index.ts @@ -1 +1,4 @@ import { Logger } from "@trigger.dev/core/logger"; + +export { TQuery } from "./query/TQuery.js"; +export type { TQueryOptions } from "./query/TQuery.js"; diff --git a/internal-packages/tsql/src/query/ClickHouseQueryVisitor.ts b/internal-packages/tsql/src/query/ClickHouseQueryVisitor.ts new file mode 100644 index 000000000..d4749315b --- /dev/null +++ b/internal-packages/tsql/src/query/ClickHouseQueryVisitor.ts @@ -0,0 +1,314 @@ +import { ParserRuleContext } from "antlr4ts/ParserRuleContext"; +import { ErrorNode } from "antlr4ts/tree/ErrorNode"; +import { ParseTree } from "antlr4ts/tree/ParseTree"; +import { TerminalNode } from "antlr4ts/tree/TerminalNode"; +import { + ArrayJoinClauseContext, + ColumnExprContext, + ColumnExprListContext, + FromClauseContext, + GroupByClauseContext, + HavingClauseContext, + JoinExprContext, + LimitAndOffsetClauseContext, + LimitByClauseContext, + OffsetOnlyClauseContext, + OrderByClauseContext, + PlaceholderContext, + PrewhereClauseContext, + SelectContext, + SelectSetStmtContext, + SelectStmtContext, + SelectStmtWithParensContext, + SettingsClauseContext, + TopClauseContext, + WhereClauseContext, + WindowClauseContext, + WithClauseContext, +} from "../grammar/TSQLParser.js"; +import { TSQLParserVisitor } from "../grammar/TSQLParserVisitor.js"; + +/** + * Visitor that converts TSQL AST to ClickHouse SQL + */ +export class ClickHouseQueryVisitor implements TSQLParserVisitor { + visitSelect(ctx: SelectContext): string { + const selectSetStmt = ctx.selectSetStmt(); + if (selectSetStmt) { + return this.visitSelectSetStmt(selectSetStmt); + } + + const selectStmt = ctx.selectStmt(); + if (selectStmt) { + return this.visitSelectStmt(selectStmt); + } + + // Handle tSQLxTagElement if needed + return ""; + } + + visitSelectSetStmt(ctx: SelectSetStmtContext): string { + const selectStmtWithParens = ctx.selectStmtWithParens(); + let query = this.visitSelectStmtWithParens(selectStmtWithParens); + + // Handle subsequent select set clauses (UNION, EXCEPT, INTERSECT) + const subsequentClauses = ctx.subsequentSelectSetClause(); + for (const clause of subsequentClauses) { + const op = clause.EXCEPT() + ? "EXCEPT" + : clause.UNION() + ? clause.ALL() + ? "UNION ALL" + : clause.DISTINCT() + ? "UNION DISTINCT" + : "UNION" + : clause.INTERSECT() + ? clause.DISTINCT() + ? "INTERSECT DISTINCT" + : "INTERSECT" + : ""; + + if (op) { + const nextSelect = clause.selectStmtWithParens(); + query += ` ${op} ${this.visitSelectStmtWithParens(nextSelect)}`; + } + } + + return query; + } + + visitSelectStmtWithParens(ctx: SelectStmtWithParensContext): string { + const selectStmt = ctx.selectStmt(); + if (selectStmt) { + return this.visitSelectStmt(selectStmt); + } + + const selectSetStmt = ctx.selectSetStmt(); + if (selectSetStmt) { + return `(${this.visitSelectSetStmt(selectSetStmt)})`; + } + + // Handle placeholder if needed + const placeholder = ctx.placeholder(); + if (placeholder) { + return this.visitPlaceholder(placeholder); + } + + return this.getTextFromContext(ctx); + } + + visitPlaceholder(ctx: PlaceholderContext): string { + return this.getTextFromContext(ctx); + } + + visitSelectStmt(ctx: SelectStmtContext): string { + let parts: string[] = []; + + // WITH clause + const withClause = ctx.withClause(); + if (withClause) { + parts.push(this.visitWithClause(withClause)); + } + + // SELECT + parts.push("SELECT"); + + // DISTINCT + if (ctx.DISTINCT()) { + parts.push("DISTINCT"); + } + + // TOP clause + const topClause = ctx.topClause(); + if (topClause) { + parts.push(this.visitTopClause(topClause)); + } + + // Column list + const columnExprList = ctx.columnExprList(); + if (columnExprList) { + parts.push(this.visitColumnExprList(columnExprList)); + } + + // FROM clause + const fromClause = ctx.fromClause(); + if (fromClause) { + parts.push(this.visitFromClause(fromClause)); + } + + // Array JOIN + const arrayJoinClause = ctx.arrayJoinClause(); + if (arrayJoinClause) { + parts.push(this.visitArrayJoinClause(arrayJoinClause)); + } + + // PREWHERE + const prewhereClause = ctx.prewhereClause(); + if (prewhereClause) { + parts.push(this.visitPrewhereClause(prewhereClause)); + } + + // WHERE + const whereClause = ctx.whereClause(); + if (whereClause) { + parts.push(this.visitWhereClause(whereClause)); + } + + // GROUP BY + const groupByClause = ctx.groupByClause(); + if (groupByClause) { + parts.push(this.visitGroupByClause(groupByClause)); + } + + // HAVING + const havingClause = ctx.havingClause(); + if (havingClause) { + parts.push(this.visitHavingClause(havingClause)); + } + + // WINDOW + const windowClause = ctx.windowClause(); + if (windowClause) { + parts.push(this.visitWindowClause(windowClause)); + } + + // ORDER BY + const orderByClause = ctx.orderByClause(); + if (orderByClause) { + parts.push(this.visitOrderByClause(orderByClause)); + } + + // LIMIT BY + const limitByClause = ctx.limitByClause(); + if (limitByClause) { + parts.push(this.visitLimitByClause(limitByClause)); + } + + // LIMIT / OFFSET + const limitAndOffsetClause = ctx.limitAndOffsetClause(); + if (limitAndOffsetClause) { + parts.push(this.visitLimitAndOffsetClause(limitAndOffsetClause)); + } + + const offsetOnlyClause = ctx.offsetOnlyClause(); + if (offsetOnlyClause) { + parts.push(this.visitOffsetOnlyClause(offsetOnlyClause)); + } + + // SETTINGS + const settingsClause = ctx.settingsClause(); + if (settingsClause) { + parts.push(this.visitSettingsClause(settingsClause)); + } + + return parts.join(" "); + } + + visitColumnExprList(ctx: ColumnExprListContext): string { + const exprs: string[] = []; + const columnExprs = ctx.columnExpr(); + for (const expr of columnExprs) { + exprs.push(this.visitColumnExpr(expr)); + } + return exprs.join(", "); + } + + visitColumnExpr(ctx: ColumnExprContext): string { + // Use the text property which contains the original input text for this node + return this.getTextFromContext(ctx); + } + + visitFromClause(ctx: FromClauseContext): string { + const joinExpr = ctx.joinExpr(); + return `FROM ${this.visitJoinExpr(joinExpr)}`; + } + + visitJoinExpr(ctx: JoinExprContext): string { + return this.getTextFromContext(ctx); + } + + visitWhereClause(ctx: WhereClauseContext): string { + const columnExpr = ctx.columnExpr(); + return `WHERE ${this.visitColumnExpr(columnExpr)}`; + } + + visitGroupByClause(ctx: GroupByClauseContext): string { + return this.getTextFromContext(ctx); + } + + visitHavingClause(ctx: HavingClauseContext): string { + return this.getTextFromContext(ctx); + } + + visitOrderByClause(ctx: OrderByClauseContext): string { + return this.getTextFromContext(ctx); + } + + visitLimitByClause(ctx: LimitByClauseContext): string { + return this.getTextFromContext(ctx); + } + + visitLimitAndOffsetClause(ctx: LimitAndOffsetClauseContext): string { + return this.getTextFromContext(ctx); + } + + visitOffsetOnlyClause(ctx: OffsetOnlyClauseContext): string { + return this.getTextFromContext(ctx); + } + + visitSettingsClause(ctx: SettingsClauseContext): string { + return this.getTextFromContext(ctx); + } + + visitWithClause(ctx: WithClauseContext): string { + return this.getTextFromContext(ctx); + } + + visitTopClause(ctx: TopClauseContext): string { + return this.getTextFromContext(ctx); + } + + visitArrayJoinClause(ctx: ArrayJoinClauseContext): string { + return this.getTextFromContext(ctx); + } + + visitPrewhereClause(ctx: PrewhereClauseContext): string { + return this.getTextFromContext(ctx); + } + + visitWindowClause(ctx: WindowClauseContext): string { + return this.getTextFromContext(ctx); + } + + /** + * Extract the original text from a parse tree context + * This uses the text property which contains the original input text + */ + private getTextFromContext(ctx: ParserRuleContext): string { + // The text property exists at runtime but may not be in types + return (ctx as any).text || ""; + } + + // Required by ParseTreeVisitor interface + visit(tree: ParseTree): string { + // Use the text property which contains the original input text + return (tree as any).text || ""; + } + + visitChildren(node: ParserRuleContext): string { + // Visit all children and concatenate their results + if (!node.children || node.children.length === 0) { + return this.visit(node); + } + return node.children.map((child: ParseTree) => this.visit(child)).join(" "); + } + + // Visit terminal nodes + visitTerminal(node: TerminalNode): string { + return node.text; + } + + visitErrorNode(node: ErrorNode): string { + return ""; + } +} diff --git a/internal-packages/tsql/src/query/TQuery.ts b/internal-packages/tsql/src/query/TQuery.ts new file mode 100644 index 000000000..e9900aca6 --- /dev/null +++ b/internal-packages/tsql/src/query/TQuery.ts @@ -0,0 +1,122 @@ +import { CharStreams, CommonTokenStream } from "antlr4ts"; +import { TSQLLexer } from "../grammar/TSQLLexer.js"; +import { TSQLParser } from "../grammar/TSQLParser.js"; +import { ClickHouseQueryVisitor } from "./ClickHouseQueryVisitor.js"; +import type { ClickHouse } from "@internal/clickhouse"; +import { z } from "zod"; + +export interface TQueryOptions { + organizationId: string; + projectId: string; + environmentId: string; +} + +export class TQuery { + private readonly organizationId: string; + private readonly projectId: string; + private readonly environmentId: string; + private readonly clickhouseReader: ClickHouse; + + constructor(clickhouseReader: ClickHouse, options: TQueryOptions) { + this.clickhouseReader = clickhouseReader; + this.organizationId = options.organizationId; + this.projectId = options.projectId; + this.environmentId = options.environmentId; + } + + /** + * Execute a TSQL query and return the results + * @param input TSQL query string + * @param schema Zod schema for the output rows + * @returns Promise with query results + */ + async query>( + input: string, + schema: TOutput + ): Promise<[Error | null, z.output[] | null]> { + // Parse the TSQL input + const inputStream = CharStreams.fromString(input); + const lexer = new TSQLLexer(inputStream); + const tokenStream = new CommonTokenStream(lexer as any); + const parser = new TSQLParser(tokenStream); + + // Parse as a SELECT statement + const tree = parser.select(); + + // Convert AST to ClickHouse SQL + const visitor = new ClickHouseQueryVisitor(); + let clickhouseQuery = visitor.visit(tree); + + // Add WHERE clauses for scoping + clickhouseQuery = this.addScopingWhereClauses(clickhouseQuery); + + // Execute the query using ClickHouse client + const queryFunction = this.clickhouseReader.reader.query({ + name: "tsql-query", + query: clickhouseQuery, + params: z.object({ + organizationId: z.string(), + projectId: z.string(), + environmentId: z.string(), + }), + schema, + }); + + return await queryFunction({ + organizationId: this.organizationId, + projectId: this.projectId, + environmentId: this.environmentId, + }); + } + + /** + * Add WHERE clauses for organization_id, project_id, and environment_id + * If the query already has a WHERE clause, we add AND conditions + */ + private addScopingWhereClauses(query: string): string { + const scopingConditions = [ + "organization_id = {organizationId: String}", + "project_id = {projectId: String}", + "environment_id = {environmentId: String}", + ].join(" AND "); + + const upperQuery = query.toUpperCase(); + + // Check if query already has a WHERE clause + const whereIndex = upperQuery.indexOf(" WHERE "); + if (whereIndex !== -1) { + // Find the end of the WHERE clause (before GROUP BY, HAVING, ORDER BY, LIMIT, etc.) + const groupByIndex = upperQuery.indexOf(" GROUP BY ", whereIndex); + const havingIndex = upperQuery.indexOf(" HAVING ", whereIndex); + const orderByIndex = upperQuery.indexOf(" ORDER BY ", whereIndex); + const limitIndex = upperQuery.indexOf(" LIMIT ", whereIndex); + + let whereEndIndex = query.length; + if (groupByIndex !== -1) whereEndIndex = Math.min(whereEndIndex, groupByIndex); + if (havingIndex !== -1) whereEndIndex = Math.min(whereEndIndex, havingIndex); + if (orderByIndex !== -1) whereEndIndex = Math.min(whereEndIndex, orderByIndex); + if (limitIndex !== -1) whereEndIndex = Math.min(whereEndIndex, limitIndex); + + // Insert AND conditions before the end of WHERE clause + const beforeWhereEnd = query.substring(0, whereEndIndex); + const afterWhereEnd = query.substring(whereEndIndex); + return `${beforeWhereEnd} AND ${scopingConditions}${afterWhereEnd}`; + } else { + // Add WHERE clause before GROUP BY, HAVING, ORDER BY, LIMIT, etc. + const groupByIndex = upperQuery.indexOf(" GROUP BY "); + const havingIndex = upperQuery.indexOf(" HAVING "); + const orderByIndex = upperQuery.indexOf(" ORDER BY "); + const limitIndex = upperQuery.indexOf(" LIMIT "); + + let insertIndex = query.length; + if (groupByIndex !== -1) insertIndex = Math.min(insertIndex, groupByIndex); + if (havingIndex !== -1) insertIndex = Math.min(insertIndex, havingIndex); + if (orderByIndex !== -1) insertIndex = Math.min(insertIndex, orderByIndex); + if (limitIndex !== -1) insertIndex = Math.min(insertIndex, limitIndex); + + return `${query.substring(0, insertIndex)} WHERE ${scopingConditions}${query.substring( + insertIndex + )}`; + } + } +} diff --git a/internal-packages/tsql/tsconfig.json b/internal-packages/tsql/tsconfig.json index 010433962..98c833a18 100644 --- a/internal-packages/tsql/tsconfig.json +++ b/internal-packages/tsql/tsconfig.json @@ -2,8 +2,8 @@ "compilerOptions": { "target": "ES2019", "lib": ["ES2019", "DOM", "DOM.Iterable", "DOM.AsyncIterable"], - "module": "Node16", - "moduleResolution": "Node16", + "module": "ESNext", + "moduleResolution": "node", "moduleDetection": "force", "verbatimModuleSyntax": false, "types": ["vitest/globals"], @@ -16,7 +16,9 @@ "strict": true, "paths": { "@trigger.dev/core": ["../../packages/core/src/index"], - "@trigger.dev/core/*": ["../../packages/core/src/*"] + "@trigger.dev/core/*": ["../../packages/core/src/*"], + "@internal/clickhouse": ["../clickhouse/src/index"], + "@internal/clickhouse/*": ["../clickhouse/src/*"] } }, "exclude": ["node_modules"] diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 09fcb2c67..cba37a57c 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -1297,12 +1297,18 @@ importers: internal-packages/tsql: dependencies: + '@internal/clickhouse': + specifier: workspace:* + version: link:../clickhouse '@trigger.dev/core': specifier: workspace:* version: link:../../packages/core antlr4ts: specifier: 0.5.0-alpha.4 version: 0.5.0-alpha.4 + zod: + specifier: 3.25.76 + version: 3.25.76 devDependencies: antlr4ts-cli: specifier: 0.5.0-alpha.4