From bdd95cf5e359a4e09df840b6c018d0566127b74e Mon Sep 17 00:00:00 2001
From: Rafael Bey <24432403+rafaelbey@users.noreply.github.com>
Date: Thu, 19 Oct 2023 17:48:06 -0400
Subject: [PATCH] Allow parsing sql expressions (#2389)

---
 .../sql/grammar/from/SQLGrammarParser.java    |  7 ++++
 .../test/roundtrip/TestSQLRoundTrip.java      | 42 +++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/legend-engine-xts-sql/legend-engine-xt-sql-grammar/src/main/java/org/finos/legend/engine/language/sql/grammar/from/SQLGrammarParser.java b/legend-engine-xts-sql/legend-engine-xt-sql-grammar/src/main/java/org/finos/legend/engine/language/sql/grammar/from/SQLGrammarParser.java
index 408f72c14c9..83be78cddba 100644
--- a/legend-engine-xts-sql/legend-engine-xt-sql-grammar/src/main/java/org/finos/legend/engine/language/sql/grammar/from/SQLGrammarParser.java
+++ b/legend-engine-xts-sql/legend-engine-xt-sql-grammar/src/main/java/org/finos/legend/engine/language/sql/grammar/from/SQLGrammarParser.java
@@ -28,6 +28,7 @@
 import org.finos.legend.engine.language.sql.grammar.from.antlr4.SqlBaseLexer;
 import org.finos.legend.engine.language.sql.grammar.from.antlr4.SqlBaseParser;
 import org.finos.legend.engine.protocol.pure.v1.model.SourceInformation;
+import org.finos.legend.engine.protocol.sql.metamodel.Expression;
 import org.finos.legend.engine.protocol.sql.metamodel.Statement;
 
 import java.util.BitSet;
@@ -51,6 +52,12 @@ public Statement parseStatement(String query)
         return this.parse(query, "statement");
     }
 
+    public Expression parseExpression(String expression)
+    {
+        SqlBaseParser parser = getSqlBaseParser(expression, "expression");
+        return (Expression) sqlVisitor.visitSingleExpression(parser.singleExpression());
+    }
+
     private Statement parse(String query, String name)
     {
         SqlBaseParser parser = getSqlBaseParser(query, name);
diff --git a/legend-engine-xts-sql/legend-engine-xt-sql-grammar/src/test/java/org/finos/legend/engine/language/sql/grammar/test/roundtrip/TestSQLRoundTrip.java b/legend-engine-xts-sql/legend-engine-xt-sql-grammar/src/test/java/org/finos/legend/engine/language/sql/grammar/test/roundtrip/TestSQLRoundTrip.java
index befce651ac1..2643bb292bc 100644
--- a/legend-engine-xts-sql/legend-engine-xt-sql-grammar/src/test/java/org/finos/legend/engine/language/sql/grammar/test/roundtrip/TestSQLRoundTrip.java
+++ b/legend-engine-xts-sql/legend-engine-xt-sql-grammar/src/test/java/org/finos/legend/engine/language/sql/grammar/test/roundtrip/TestSQLRoundTrip.java
@@ -128,18 +128,36 @@ public void testWhere()
         check("SELECT * FROM myTable WHERE col1 = 1");
     }
 
+    @Test
+    public void testWhereExpression()
+    {
+        checkExpression("col1 = 1");
+    }
+
     @Test
     public void testCompositeWhere()
     {
         check("SELECT * FROM myTable WHERE col1 = 1 AND col2 = 1");
     }
 
+    @Test
+    public void testCompositeWhereExpression()
+    {
+        checkExpression("col1 = 1 AND col2 = 1");
+    }
+
     @Test
     public void testWhereQualified()
     {
         check("SELECT * FROM myTable WHERE myTable.col1 = 1");
     }
 
+    @Test
+    public void testWhereQualifiedExpression()
+    {
+        checkExpression("myTable.col1 = 1");
+    }
+
     @Test
     public void testCompositeWhereQualifiedWithAlias()
     {
@@ -155,6 +173,15 @@ public void testCompositeWhereOperators()
                 "col BETWEEN 0 AND 1");
     }
 
+    @Test
+    public void testCompositeWhereOperatorsExpression()
+    {
+        checkExpression("col = 1 AND col > 1 AND col < 1 " +
+                "AND col >= 1 AND col <= 1 AND col IN (1, 2, 3) AND col IS NULL AND " +
+                "col IS NOT NULL AND col IS DISTINCT FROM 1 AND col IS NOT DISTINCT FROM 1 AND " +
+                "col BETWEEN 0 AND 1");
+    }
+
     @Test
     public void testGroupBy()
     {
@@ -324,4 +351,19 @@ private void check(String sql, String expected)
         String result = composer.renderNode(node);
         MatcherAssert.assertThat(result.trim(), IsEqualIgnoringCase.equalToIgnoringCase(expected));
     }
+
+    private void checkExpression(String expression)
+    {
+        checkExpression(expression, expression);
+        checkExpression(expression.toLowerCase(), expression);
+    }
+
+    private void checkExpression(String expression, String expected)
+    {
+        SQLGrammarParser parser = SQLGrammarParser.newInstance();
+        Node node = parser.parseExpression(expression);
+        SQLGrammarComposer composer = SQLGrammarComposer.newInstance();
+        String result = composer.renderNode(node);
+        MatcherAssert.assertThat(result.trim(), IsEqualIgnoringCase.equalToIgnoringCase(expected));
+    }
 }