refactor

linkedin · Nov 7, 2023 · 833f139 · 833f139
1 parent a37ccc0
commit 833f139
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 19 deletions.
diff --git a/coral-spark/src/main/java/com/linkedin/coral/spark/dialect/SparkSqlDialect.java b/coral-spark/src/main/java/com/linkedin/coral/spark/dialect/SparkSqlDialect.java
@@ -5,9 +5,13 @@
  */
 package com.linkedin.coral.spark.dialect;
 
-import org.apache.calcite.avatica.util.Casing;
-import com.google.common.collect.ImmutableList;
+import java.util.Arrays;
 import java.util.List;
+import java.util.Locale;
+
+import com.google.common.collect.ImmutableList;
+
+import org.apache.calcite.avatica.util.Casing;
 import org.apache.calcite.config.NullCollation;
 import org.apache.calcite.sql.SqlCall;
 import org.apache.calcite.sql.SqlDialect;
@@ -33,13 +37,26 @@
  */
 public class SparkSqlDialect extends SqlDialect {
 
-  public static final SqlDialect.Context DEFAULT_CONTEXT =
-      SqlDialect.EMPTY_CONTEXT.withDatabaseProduct(DatabaseProduct.SPARK).withLiteralQuoteString("'")
-          .withLiteralEscapedQuoteString("\\'").withNullCollation(NullCollation.LOW)
-          .withUnquotedCasing(Casing.UNCHANGED).withQuotedCasing(Casing.UNCHANGED).withCaseSensitive(false).withIdentifierQuoteString("`");
+  public static final SqlDialect.Context DEFAULT_CONTEXT = SqlDialect.EMPTY_CONTEXT
+      .withDatabaseProduct(DatabaseProduct.SPARK).withLiteralQuoteString("'").withLiteralEscapedQuoteString("\\'")
+      .withNullCollation(NullCollation.LOW).withUnquotedCasing(Casing.UNCHANGED).withQuotedCasing(Casing.UNCHANGED)
+      .withCaseSensitive(false).withIdentifierQuoteString("`");
 
   public static final SparkSqlDialect INSTANCE = new SparkSqlDialect(DEFAULT_CONTEXT);
 
+  private static final List<String> RESERVED_KEYWORDS = ImmutableList.copyOf(
+      Arrays.asList("ALL", "ALTER", "AND", "ARRAY", "AS", "AUTHORIZATION", "BETWEEN", "BIGINT", "BINARY", "BOOLEAN",
+          "BOTH", "BY", "CASE", "CAST", "CHAR", "COLUMN", "CONF", "CREATE", "CROSS", "CUBE", "CURRENT", "CURRENT_DATE",
+          "CURRENT_TIMESTAMP", "CURSOR", "DATABASE", "DATE", "DECIMAL", "DELETE", "DESCRIBE", "DISTINCT", "DOUBLE",
+          "DROP", "ELSE", "END", "EXCHANGE", "EXISTS", "EXTENDED", "EXTERNAL", "FALSE", "FETCH", "FLOAT", "FOLLOWING",
+          "FOR", "FROM", "FULL", "FUNCTION", "GRANT", "GROUP", "GROUPING", "HAVING", "IF", "IMPORT", "IN", "INNER",
+          "INSERT", "INT", "INTERSECT", "INTERVAL", "INTO", "IS", "JOIN", "LATERAL", "LEFT", "LESS", "LIKE", "LOCAL",
+          "MACRO", "MAP", "MORE", "NONE", "NOT", "NULL", "OF", "ON", "OR", "ORDER", "OUT", "OUTER", "OVER",
+          "PARTIALSCAN", "PARTITION", "PERCENT", "PRECEDING", "PRESERVE", "PROCEDURE", "RANGE", "READS", "REDUCE",
+          "REVOKE", "RIGHT", "ROLLUP", "ROW", "ROWS", "SELECT", "SET", "SMALLINT", "TABLE", "TABLESAMPLE", "THEN",
+          "TIMESTAMP", "TO", "TRANSFORM", "TRIGGER", "TRUE", "TRUNCATE", "UNBOUNDED", "UNION", "UNIQUEJOIN", "UPDATE",
+          "USER", "USING", "UTC_TMESTAMP", "VALUES", "VARCHAR", "WHEN", "WHERE", "WINDOW", "WITH"));
+
   private SparkSqlDialect(Context context) {
     super(context);
   }
@@ -137,20 +154,12 @@ public boolean supportsCharSet() {
     return false;
   }
 
-  public String quoteIdentifier(String val) {
-    List<String> reservedKeywords = ImmutableList.of("select", "timestamp");
-    if (reservedKeywords.contains(val.toLowerCase())) {
-      String val2 =
-          val.replaceAll(
-              identifierEndQuoteString,
-              identifierEscapedQuote);
-      return identifierQuoteString + val2 + identifierEndQuoteString;
-    } else {
-      return val;
-    }
+  @Override
+  protected boolean identifierNeedsQuote(String val) {
+    return RESERVED_KEYWORDS.contains(val.toUpperCase(Locale.ROOT));
   }
 
-
+  @Override
   public void unparseOffsetFetch(SqlWriter writer, SqlNode offset, SqlNode fetch) {
     unparseFetchUsingLimit(writer, offset, fetch);
   }

diff --git a/coral-spark/src/test/java/com/linkedin/coral/spark/CoralSparkTest.java b/coral-spark/src/test/java/com/linkedin/coral/spark/CoralSparkTest.java
@@ -73,7 +73,7 @@ public void testGetBaseTablesFromView() {
   }
 
   @Test
-  public void testTimestamp() {
+  public void testQuotingKeywords() {
     RelNode relNode = TestUtils.toRelNode("default", "baz_view");
     CoralSpark coralSpark = createCoralSpark(relNode);
     assertEquals(coralSpark.getSparkSql(), "SELECT baz.`select`, baz.`timestamp`\n" + "FROM default.baz baz");