Sanitize Symbol names for unambiguous plans

trinodb · Mar 22, 2024 · 5c19d1c · 5c19d1c
1 parent 84f8c05
commit 5c19d1c
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 2 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/Symbol.java b/core/trino-main/src/main/java/io/trino/sql/planner/Symbol.java
@@ -16,16 +16,25 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.databind.annotation.JsonSerialize;
+import com.google.common.base.CharMatcher;
 import io.trino.spi.type.Type;
 import io.trino.sql.ir.Expression;
 import io.trino.sql.ir.Reference;
 
+import static com.google.common.base.Preconditions.checkArgument;
 import static java.util.Objects.requireNonNull;
 
 @JsonSerialize(keyUsing = SymbolKeySerializer.class)
 public class Symbol
         implements Comparable<Symbol>
 {
+    // Allow limited set of characters in symbol names to avoid ambiguities in the rendered plans.
+    public static final CharMatcher SYMBOL_NAME_MATCHER = CharMatcher.inRange('a', 'z')
+            .or(CharMatcher.inRange('A', 'Z'))
+            .or(CharMatcher.inRange('0', '9'))
+            .or(CharMatcher.anyOf("[]:_$"))
+            .precomputed();
+
     private final String name;
     private final Type type;
 
@@ -40,8 +49,9 @@ public static Symbol from(Expression expression)
     @JsonCreator
     public Symbol(Type type, String name)
     {
-        requireNonNull(name, "name is null");
         requireNonNull(type, "type is null");
+        checkArgument(!name.isEmpty(), "name cannot be empty");
+        checkArgument(SYMBOL_NAME_MATCHER.matchesAllOf(name), "Invalid characters in symbol name: [%s]", name);
         this.type = type;
         this.name = name;
     }

diff --git a/core/trino-main/src/main/java/io/trino/sql/planner/SymbolAllocator.java b/core/trino-main/src/main/java/io/trino/sql/planner/SymbolAllocator.java
@@ -26,6 +26,8 @@
 import java.util.HashMap;
 import java.util.Map;
 
+import static com.google.common.base.Preconditions.checkArgument;
+import static io.trino.sql.planner.Symbol.SYMBOL_NAME_MATCHER;
 import static java.util.Locale.ENGLISH;
 import static java.util.Objects.requireNonNull;
 import static java.util.stream.Collectors.toMap;
@@ -85,9 +87,10 @@ public Symbol newSymbol(String nameHint, Type type, @Nullable String suffix)
             }
         }
 
-        String unique = nameHint;
+        String unique = SYMBOL_NAME_MATCHER.retainFrom(nameHint);
 
         if (suffix != null) {
+            checkArgument(SYMBOL_NAME_MATCHER.matchesAllOf(suffix), "Invalid characters in symbol suffix: [%s]", suffix);
             unique = unique + "$" + suffix;
         }