Remove identifierQuote from ConnectorExpressionRewriter

Instead, pass it to `RewriteVariable` directly, as this is the only rule needing it.
trinodb · Mar 16, 2022 · b704cee · b704cee
1 parent 8add1e7
commit b704cee
Show file tree

Hide file tree

Showing 6 changed files with 15 additions and 35 deletions.
diff --git a/...in-toolkit/src/main/java/io/trino/plugin/base/expression/ConnectorExpressionRewriter.java b/...in-toolkit/src/main/java/io/trino/plugin/base/expression/ConnectorExpressionRewriter.java
@@ -26,20 +26,17 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
-import java.util.function.Function;
 
 import static com.google.common.base.Verify.verify;
 import static io.trino.matching.Capture.newCapture;
 import static java.util.Objects.requireNonNull;
 
 public final class ConnectorExpressionRewriter<Result>
 {
-    private final Function<String, String> identifierQuote;
     private final Set<ConnectorExpressionRule<?, Result>> rules;
 
-    public ConnectorExpressionRewriter(Function<String, String> identifierQuote, Set<ConnectorExpressionRule<?, Result>> rules)
+    public ConnectorExpressionRewriter(Set<ConnectorExpressionRule<?, Result>> rules)
     {
-        this.identifierQuote = requireNonNull(identifierQuote, "identifierQuote is null");
         this.rules = ImmutableSet.copyOf(requireNonNull(rules, "rules is null"));
     }
 
@@ -57,12 +54,6 @@ public Map<String, ColumnHandle> getAssignments()
                 return assignments;
             }
 
-            @Override
-            public Function<String, String> getIdentifierQuote()
-            {
-                return identifierQuote;
-            }
-
             @Override
             public ConnectorSession getSession()
             {

diff --git a/...plugin-toolkit/src/main/java/io/trino/plugin/base/expression/ConnectorExpressionRule.java b/...plugin-toolkit/src/main/java/io/trino/plugin/base/expression/ConnectorExpressionRule.java
@@ -21,7 +21,6 @@
 
 import java.util.Map;
 import java.util.Optional;
-import java.util.function.Function;
 
 import static com.google.common.base.Verify.verifyNotNull;
 import static java.util.Objects.requireNonNull;
@@ -44,8 +43,6 @@ default ColumnHandle getAssignment(String name)
 
         Map<String, ColumnHandle> getAssignments();
 
-        Function<String, String> getIdentifierQuote();
-
         ConnectorSession getSession();
 
         Optional<Result> defaultRewrite(ConnectorExpression expression);

diff --git a/...src/main/java/io/trino/plugin/jdbc/expression/JdbcConnectorExpressionRewriterBuilder.java b/...src/main/java/io/trino/plugin/jdbc/expression/JdbcConnectorExpressionRewriterBuilder.java
@@ -19,29 +19,20 @@
 
 import java.util.function.Function;
 
-import static java.util.Objects.requireNonNull;
-
 public class JdbcConnectorExpressionRewriterBuilder
 {
     public static JdbcConnectorExpressionRewriterBuilder newBuilder()
     {
         return new JdbcConnectorExpressionRewriterBuilder();
     }
 
-    private Function<String, String> identifierQuote;
     private ImmutableSet.Builder<ConnectorExpressionRule<?, String>> rules = ImmutableSet.builder();
 
     private JdbcConnectorExpressionRewriterBuilder() {}
 
-    public JdbcConnectorExpressionRewriterBuilder withIdentifierQuote(Function<String, String> identifierQuote)
-    {
-        this.identifierQuote = requireNonNull(identifierQuote, "identifierQuote is null");
-        return this;
-    }
-
-    public JdbcConnectorExpressionRewriterBuilder addStandardRules()
+    public JdbcConnectorExpressionRewriterBuilder addStandardRules(Function<String, String> identifierQuote)
     {
-        add(new RewriteVariable());
+        add(new RewriteVariable(identifierQuote));
         add(new RewriteVarcharConstant());
         add(new RewriteExactNumericConstant());
         add(new RewriteOr());
@@ -70,7 +61,7 @@ public JdbcConnectorExpressionRewriterBuilder to(String rewritePattern)
 
     public ConnectorExpressionRewriter<String> build()
     {
-        return new ConnectorExpressionRewriter<>(this.identifierQuote, rules.build());
+        return new ConnectorExpressionRewriter<>(rules.build());
     }
 
     public interface ExpressionMapping<Continuation>

diff --git a/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/expression/RewriteVariable.java b/plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/expression/RewriteVariable.java
@@ -20,12 +20,21 @@
 import io.trino.spi.expression.Variable;
 
 import java.util.Optional;
+import java.util.function.Function;
 
 import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.variable;
+import static java.util.Objects.requireNonNull;
 
 public class RewriteVariable
         implements ConnectorExpressionRule<Variable, String>
 {
+    private final Function<String, String> identifierQuote;
+
+    public RewriteVariable(Function<String, String> identifierQuote)
+    {
+        this.identifierQuote = requireNonNull(identifierQuote, "identifierQuote is null");
+    }
+
     @Override
     public Pattern<Variable> getPattern()
     {
@@ -36,6 +45,6 @@ public Pattern<Variable> getPattern()
     public Optional<String> rewrite(Variable variable, Captures captures, RewriteContext<String> context)
     {
         JdbcColumnHandle columnHandle = (JdbcColumnHandle) context.getAssignment(variable.getName());
-        return Optional.of(context.getIdentifierQuote().apply(columnHandle.getColumnName()));
+        return Optional.of(identifierQuote.apply(columnHandle.getColumnName()));
     }
 }
diff --git a/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/expression/TestGenericRewrite.java b/plugin/trino-base-jdbc/src/test/java/io/trino/plugin/jdbc/expression/TestGenericRewrite.java
@@ -26,7 +26,6 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
-import java.util.function.Function;
 
 import static com.google.common.collect.MoreCollectors.onlyElement;
 import static io.trino.spi.type.BigintType.BIGINT;
@@ -55,12 +54,6 @@ public Map<String, ColumnHandle> getAssignments()
                 throw new UnsupportedOperationException();
             }
 
-            @Override
-            public Function<String, String> getIdentifierQuote()
-            {
-                throw new UnsupportedOperationException();
-            }
-
             @Override
             public ConnectorSession getSession()
             {

diff --git a/plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlClient.java b/plugin/trino-postgresql/src/main/java/io/trino/plugin/postgresql/PostgreSqlClient.java
@@ -304,8 +304,7 @@ public PostgreSqlClient(
                         .build());
 
         connectorExpressionRewriter = JdbcConnectorExpressionRewriterBuilder.newBuilder()
-                .withIdentifierQuote(this::quoted)
-                .addStandardRules()
+                .addStandardRules(this::quoted)
                 // TODO allow all comparison operators for numeric types
                 .add(new RewriteComparison(RewriteComparison.ComparisonOperator.EQUAL, RewriteComparison.ComparisonOperator.NOT_EQUAL))
                 .map("$like_pattern(value: varchar, pattern: varchar): boolean").to("value LIKE pattern")