DSL for JDBC expression rewrites

plugin/trino-base-jdbc/pom.xml

@@ -124,6 +124,12 @@
             <artifactId>failsafe</artifactId>
         </dependency>
 
+        <dependency>
+            <groupId>org.antlr</groupId>
+            <artifactId>antlr4-runtime</artifactId>
+            <version>${dep.antlr.version}</version>
+        </dependency>
+
         <dependency>
             <groupId>org.weakref</groupId>
             <artifactId>jmxutils</artifactId>
@@ -238,4 +244,13 @@
             <scope>test</scope>
         </dependency>
     </dependencies>
+
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>org.antlr</groupId>
+                <artifactId>antlr4-maven-plugin</artifactId>
+            </plugin>
+        </plugins>
+    </build>
 </project>

plugin/trino-base-jdbc/src/main/antlr4/io/trino/plugin/jdbc/expression/ConnectorExpressionPattern.g4

@@ -0,0 +1,81 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+grammar ConnectorExpressionPattern;
+
+tokens {
+    DELIMITER
+}
+
+standaloneExpression
+    : expression EOF
+    ;
+
+standaloneType
+    : type EOF
+    ;
+
+expression
+    : call
+    | expressionCapture
+    ;
+
+call
+    : identifier '(' expression (',' expression)* ')' (':' type)?
+    ;
+
+expressionCapture
+    : identifier ':' type
+    ;
+
+type
+    : identifier
+    | identifier '(' typeParameter (',' typeParameter)* ')'
+    ;
+
+typeParameter
+    : number
+    | identifier
+    ;
+
+identifier
+    : IDENTIFIER
+    ;
+
+number
+    : INTEGER_VALUE
+    ;
+
+IDENTIFIER
+    : (LETTER | '_' | '$') (LETTER | DIGIT | '_' | '$')*
+    ;
+
+INTEGER_VALUE
+    : DIGIT+
+    ;
+
+fragment DIGIT
+    : [0-9]
+    ;
+
+fragment LETTER
+    : [A-Za-z]
+    ;
+
+WS
+    : [ \r\n\t]+ -> channel(HIDDEN)
+    ;
+
+// Catch-all for anything we can't recognize.
+UNRECOGNIZED: .;

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/expression/CallPattern.java

@@ -0,0 +1,107 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.jdbc.expression;
+
+import com.google.common.collect.ImmutableList;
+import io.trino.matching.Captures;
+import io.trino.matching.Pattern;
+import io.trino.spi.expression.Call;
+import io.trino.spi.expression.ConnectorExpression;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.Optional;
+
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.argument;
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.argumentCount;
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.call;
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.functionUnqualifiedName;
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.type;
+import static java.lang.String.format;
+import static java.util.Objects.requireNonNull;
+import static java.util.stream.Collectors.joining;
+
+public class CallPattern
+        extends ExpressionPattern
+{
+    private final String functionName;
+    private final List<ExpressionPattern> parameters;
+    private final Optional<TypePattern> resultType;
+    private final Pattern<Call> pattern;
+
+    public CallPattern(String functionName, List<ExpressionPattern> parameters, Optional<TypePattern> resultType)
+    {
+        this.functionName = requireNonNull(functionName, "functionName is null");
+        this.parameters = ImmutableList.copyOf(requireNonNull(parameters, "parameters is null"));
+        this.resultType = requireNonNull(resultType, "resultType is null");
+
+        Pattern<Call> pattern = call().with(functionUnqualifiedName().equalTo(functionName));
+        if (resultType.isPresent()) {
+            pattern = pattern.with(type().matching(resultType.get().getPattern()));
+        }
+        pattern = pattern.with(argumentCount().equalTo(parameters.size()));
+        for (int i = 0; i < parameters.size(); i++) {
+            pattern = pattern.with(argument(i).matching(parameters.get(i).getPattern()));
+        }
+        this.pattern = pattern;
+    }
+
+    @Override
+    public Pattern<? extends ConnectorExpression> getPattern()
+    {
+        return pattern;
+    }
+
+    @Override
+    public void resolve(Captures captures, MatchContext matchContext)
+    {
+        for (ExpressionPattern parameter : parameters) {
+            parameter.resolve(captures, matchContext);
+        }
+        resultType.ifPresent(resultType -> resultType.resolve(captures, matchContext));
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        CallPattern that = (CallPattern) o;
+        return Objects.equals(functionName, that.functionName) &&
+                Objects.equals(parameters, that.parameters) &&
+                Objects.equals(resultType, that.resultType);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(functionName, parameters, resultType);
+    }
+
+    @Override
+    public String toString()
+    {
+        return format(
+                "%s(%s)%s",
+                functionName,
+                parameters.stream()
+                        .map(Object::toString)
+                        .collect(joining(", ")),
+                resultType.map(resultType -> ": " + resultType).orElse(""));
+    }
+}

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/expression/ExpressionCapture.java

@@ -0,0 +1,83 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.jdbc.expression;
+
+import io.trino.matching.Capture;
+import io.trino.matching.Captures;
+import io.trino.matching.Pattern;
+import io.trino.spi.expression.ConnectorExpression;
+
+import java.util.Objects;
+
+import static io.trino.matching.Capture.newCapture;
+import static io.trino.plugin.base.expression.ConnectorExpressionPatterns.type;
+import static java.lang.String.format;
+import static java.util.Objects.requireNonNull;
+
+public class ExpressionCapture
+        extends ExpressionPattern
+{
+    private final String name;
+    private final TypePattern type;
+
+    private final Capture<ConnectorExpression> capture = newCapture();
+    private final Pattern<ConnectorExpression> pattern;
+
+    public ExpressionCapture(String name, TypePattern type)
+    {
+        this.name = requireNonNull(name, "name is null");
+        this.type = requireNonNull(type, "type is null");
+        this.pattern = Pattern.typeOf(ConnectorExpression.class).capturedAs(capture)
+                .with(type().matching(type.getPattern()));
+    }
+
+    @Override
+    public Pattern<ConnectorExpression> getPattern()
+    {
+        return pattern;
+    }
+
+    @Override
+    public void resolve(Captures captures, MatchContext matchContext)
+    {
+        matchContext.record(name, captures.get(capture));
+        type.resolve(captures, matchContext);
+    }
+
+    @Override
+    public boolean equals(Object o)
+    {
+        if (this == o) {
+            return true;
+        }
+        if (o == null || getClass() != o.getClass()) {
+            return false;
+        }
+        ExpressionCapture that = (ExpressionCapture) o;
+        return Objects.equals(name, that.name) &&
+                Objects.equals(type, that.type);
+    }
+
+    @Override
+    public int hashCode()
+    {
+        return Objects.hash(name, type);
+    }
+
+    @Override
+    public String toString()
+    {
+        return format("%s: %s", name, type);
+    }
+}

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/expression/ExpressionMappingParser.java

@@ -0,0 +1,83 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.trino.plugin.jdbc.expression;
+
+import org.antlr.v4.runtime.BaseErrorListener;
+import org.antlr.v4.runtime.CharStreams;
+import org.antlr.v4.runtime.CommonTokenStream;
+import org.antlr.v4.runtime.ParserRuleContext;
+import org.antlr.v4.runtime.RecognitionException;
+import org.antlr.v4.runtime.Recognizer;
+import org.antlr.v4.runtime.atn.PredictionMode;
+import org.antlr.v4.runtime.misc.ParseCancellationException;
+
+import java.util.function.Function;
+
+import static java.lang.String.format;
+
+public class ExpressionMappingParser
+{
+    private static final BaseErrorListener ERROR_LISTENER = new BaseErrorListener()
+    {
+        @Override
+        public void syntaxError(Recognizer<?, ?> recognizer, Object offendingSymbol, int line, int charPositionInLine, String message, RecognitionException e)
+        {
+            throw new IllegalArgumentException(format("Error at %s:%s: %s", line, charPositionInLine, message), e);
+        }
+    };
+
+    public ExpressionPattern createExpressionPattern(String expressionPattern)
+    {
+        return (ExpressionPattern) invokeParser(expressionPattern, ConnectorExpressionPatternParser::standaloneExpression);
+    }
+
+    public TypePattern createTypePattern(String typePattern)
+    {
+        return (TypePattern) invokeParser(typePattern, ConnectorExpressionPatternParser::standaloneType);
+    }
+
+    public Object invokeParser(String input, Function<ConnectorExpressionPatternParser, ParserRuleContext> parseFunction)
+    {
+        try {
+            ConnectorExpressionPatternLexer lexer = new ConnectorExpressionPatternLexer(CharStreams.fromString(input));
+            CommonTokenStream tokenStream = new CommonTokenStream(lexer);
+            ConnectorExpressionPatternParser parser = new ConnectorExpressionPatternParser(tokenStream);
+
+            lexer.removeErrorListeners();
+            lexer.addErrorListener(ERROR_LISTENER);
+
+            parser.removeErrorListeners();
+            parser.addErrorListener(ERROR_LISTENER);
+
+            ParserRuleContext tree;
+            try {
+                // first, try parsing with potentially faster SLL mode
+                parser.getInterpreter().setPredictionMode(PredictionMode.SLL);
+                tree = parseFunction.apply(parser);
+            }
+            catch (ParseCancellationException ex) {
+                // if we fail, parse with LL mode
+                tokenStream.seek(0); // rewind input stream
+                parser.reset();
+
+                parser.getInterpreter().setPredictionMode(PredictionMode.LL);
+                tree = parseFunction.apply(parser);
+            }
+            return new ExpressionPatternBuilder().visit(tree);
+        }
+        catch (StackOverflowError e) {
+            throw new IllegalArgumentException("expression pattern is too large (stack overflow while parsing)");
+        }
+    }
+}