From c101c792df960a822346dc7a3c39f17be5cb26ec Mon Sep 17 00:00:00 2001
From: Brian Tu <brianptu2@gmail.com>
Date: Thu, 21 Nov 2024 15:48:47 -0500
Subject: [PATCH] Add custom ANTLR error listener on frontend (#26014)

## Summary & Motivation
The error message that ANTLR provides is too verbose for the user. We want to provide a more friendly error message.

## How I Tested These Changes
`AntlrAssetSelection.test.ts`
---
 .../asset-selection/AntlrAssetSelection.ts    | 49 ++++++++++++++++---
 .../__tests__/AntlrAssetSelection.test.ts     | 21 +++++++-
 2 files changed, 61 insertions(+), 9 deletions(-)

diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-selection/AntlrAssetSelection.ts b/js_modules/dagster-ui/packages/ui-core/src/asset-selection/AntlrAssetSelection.ts
index 4fc2a878c1987..fb52dc47984aa 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/asset-selection/AntlrAssetSelection.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/asset-selection/AntlrAssetSelection.ts
@@ -1,19 +1,52 @@
-import {CharStreams, CommonTokenStream} from 'antlr4ts';
+import {
+  ANTLRErrorListener,
+  CharStreams,
+  CommonTokenStream,
+  RecognitionException,
+  Recognizer,
+} from 'antlr4ts';
 
 import {AntlrAssetSelectionVisitor} from './AntlrAssetSelectionVisitor';
 import {AssetGraphQueryItem} from '../asset-graph/useAssetGraphData';
 import {AssetSelectionLexer} from './generated/AssetSelectionLexer';
 import {AssetSelectionParser} from './generated/AssetSelectionParser';
 
+class AntlrInputErrorListener implements ANTLRErrorListener<any> {
+  syntaxError(
+    recognizer: Recognizer<any, any>,
+    offendingSymbol: any,
+    line: number,
+    charPositionInLine: number,
+    msg: string,
+    e: RecognitionException | undefined,
+  ): void {
+    if (offendingSymbol) {
+      throw new Error(`Syntax error caused by "${offendingSymbol.text}": ${msg}`);
+    }
+    throw new Error(`Syntax error at char ${charPositionInLine}: ${msg}`);
+  }
+}
+
 export const parseAssetSelectionQuery = (
   all_assets: AssetGraphQueryItem[],
   query: string,
-): AssetGraphQueryItem[] => {
-  const lexer = new AssetSelectionLexer(CharStreams.fromString(query));
-  const tokenStream = new CommonTokenStream(lexer);
-  const parser = new AssetSelectionParser(tokenStream);
-  const tree = parser.start();
+): AssetGraphQueryItem[] | Error => {
+  try {
+    const lexer = new AssetSelectionLexer(CharStreams.fromString(query));
+    lexer.removeErrorListeners();
+    lexer.addErrorListener(new AntlrInputErrorListener());
+
+    const tokenStream = new CommonTokenStream(lexer);
+
+    const parser = new AssetSelectionParser(tokenStream);
+    parser.removeErrorListeners();
+    parser.addErrorListener(new AntlrInputErrorListener());
+
+    const tree = parser.start();
 
-  const visitor = new AntlrAssetSelectionVisitor(all_assets);
-  return [...visitor.visit(tree)];
+    const visitor = new AntlrAssetSelectionVisitor(all_assets);
+    return [...visitor.visit(tree)];
+  } catch (e) {
+    return e as Error;
+  }
 };
diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-selection/__tests__/AntlrAssetSelection.test.ts b/js_modules/dagster-ui/packages/ui-core/src/asset-selection/__tests__/AntlrAssetSelection.test.ts
index b484b7beb9a49..f2b2c6eed0810 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/asset-selection/__tests__/AntlrAssetSelection.test.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/asset-selection/__tests__/AntlrAssetSelection.test.ts
@@ -48,11 +48,30 @@ const TEST_GRAPH: AssetGraphQueryItem[] = [
 
 function assertQueryResult(query: string, expectedNames: string[]) {
   const result = parseAssetSelectionQuery(TEST_GRAPH, query);
+  expect(result).not.toBeInstanceOf(Error);
+  if (result instanceof Error) {
+    throw result;
+  }
   expect(result.length).toBe(expectedNames.length);
-  expect(new Set(result.map((r) => r.name))).toEqual(new Set(expectedNames));
+  expect(new Set(result.map((asset) => asset.name))).toEqual(new Set(expectedNames));
 }
 
 describe('parseAssetSelectionQuery', () => {
+  describe('invalid queries', () => {
+    it('should throw on invalid queries', () => {
+      expect(parseAssetSelectionQuery(TEST_GRAPH, 'A')).toBeInstanceOf(Error);
+      expect(parseAssetSelectionQuery(TEST_GRAPH, 'key:A key:B')).toBeInstanceOf(Error);
+      expect(parseAssetSelectionQuery(TEST_GRAPH, 'not')).toBeInstanceOf(Error);
+      expect(parseAssetSelectionQuery(TEST_GRAPH, 'and')).toBeInstanceOf(Error);
+      expect(parseAssetSelectionQuery(TEST_GRAPH, 'key:A and')).toBeInstanceOf(Error);
+      expect(parseAssetSelectionQuery(TEST_GRAPH, 'sinks')).toBeInstanceOf(Error);
+      expect(parseAssetSelectionQuery(TEST_GRAPH, 'notafunction()')).toBeInstanceOf(Error);
+      expect(parseAssetSelectionQuery(TEST_GRAPH, 'tag:foo=')).toBeInstanceOf(Error);
+      expect(parseAssetSelectionQuery(TEST_GRAPH, 'owner')).toBeInstanceOf(Error);
+      expect(parseAssetSelectionQuery(TEST_GRAPH, 'owner:owner@owner.com')).toBeInstanceOf(Error);
+    });
+  });
+
   describe('valid queries', () => {
     it('should parse star query', () => {
       assertQueryResult('*', ['A', 'B', 'C']);