Label BigQuery jobs with Trino query id

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/logging/FormatInterpolator.java

@@ -0,0 +1,89 @@
+/*
+ * 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.base.logging;
+
+import com.google.common.collect.ImmutableList;
+
+import java.util.Arrays;
+import java.util.List;
+import java.util.regex.MatchResult;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+import static com.google.common.base.MoreObjects.firstNonNull;
+import static java.util.Objects.requireNonNull;
+import static java.util.stream.Collectors.joining;
+
+public class FormatInterpolator<Context>
+{
+    private final String format;
+    private final List<InterpolatedValue<Context>> values;
+
+    public FormatInterpolator(String format, List<InterpolatedValue<Context>> values)
+    {
+        this.format = firstNonNull(format, "");
+        this.values = ImmutableList.copyOf(requireNonNull(values, "values is null"));
+    }
+
+    @SafeVarargs
+    public FormatInterpolator(String format, InterpolatedValue<Context>... values)
+    {
+        this(format, Arrays.stream(values).toList());
+    }
+
+    public String interpolate(Context context)
+    {
+        String result = format;
+        for (InterpolatedValue<Context> value : values) {
+            if (result.contains(value.getCode())) {
+                result = result.replaceAll(value.getMatchCase(), value.value(context));
+            }
+        }
+        return result;
+    }
+
+    public static boolean hasValidPlaceholders(String format, InterpolatedValue<?>... values)
+    {
+        return hasValidPlaceholders(format, Arrays.stream(values).toList());
+    }
+
+    public static boolean hasValidPlaceholders(String format, List<InterpolatedValue<?>> values)
+    {
+        List<String> matches = values.stream().map(InterpolatedValue::getMatchCase).toList();
+        Pattern pattern = Pattern.compile("[\\w ,_\\-=]|" + String.join("|", matches));
+
+        Matcher matcher = pattern.matcher(format);
+        return matcher.results()
+                .map(MatchResult::group)
+                .collect(joining())
+                .equals(format);
+    }
+
+    interface InterpolatedValue<Context>
+    {
+        String name();
+
+        default String getCode()
+        {
+            return "$" + this.name();
+        }
+
+        default String getMatchCase()
+        {
+            return "\\$" + this.name();
+        }
+
+        String value(Context context);
+    }
+}

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/logging/SessionInterpolatedValues.java

@@ -0,0 +1,69 @@
+/*
+ * 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.base.logging;
+
+import io.trino.plugin.base.logging.FormatInterpolator.InterpolatedValue;
+import io.trino.spi.connector.ConnectorSession;
+
+import java.util.function.Function;
+import java.util.function.Predicate;
+import java.util.regex.Pattern;
+
+import static java.util.Objects.requireNonNull;
+
+public enum SessionInterpolatedValues
+        implements InterpolatedValue<ConnectorSession>
+{
+    QUERY_ID(ConnectorSession::getQueryId),
+    SOURCE(new SanitizedValuesProvider(session -> session.getSource().orElse(""), "$SOURCE")),
+    USER(ConnectorSession::getUser),
+    TRACE_TOKEN(new SanitizedValuesProvider(session -> session.getTraceToken().orElse(""), "$TRACE_TOKEN"));
+
+    private final Function<ConnectorSession, String> valueProvider;
+
+    SessionInterpolatedValues(Function<ConnectorSession, String> valueProvider)
+    {
+        this.valueProvider = valueProvider;
+    }
+
+    @Override
+    public String value(ConnectorSession session)
+    {
+        return valueProvider.apply(session);
+    }
+
+    static class SanitizedValuesProvider
+            implements Function<ConnectorSession, String>
+    {
+        private static final Predicate<String> VALIDATION_MATCHER = Pattern.compile("^[\\w_-]*$").asMatchPredicate();
+        private final Function<ConnectorSession, String> valueProvider;
+        private final String name;
+
+        private SanitizedValuesProvider(Function<ConnectorSession, String> valueProvider, String name)
+        {
+            this.valueProvider = requireNonNull(valueProvider, "valueProvider is null");
+            this.name = requireNonNull(name, "name is null");
+        }
+
+        @Override
+        public String apply(ConnectorSession session)
+        {
+            String value = valueProvider.apply(session);
+            if (VALIDATION_MATCHER.test(value)) {
+                return value;
+            }
+            throw new SecurityException("Passed value %s as %s does not meet security criteria. It can contain only letters, digits, underscores and hyphens".formatted(value, name));
+        }
+    }
+}

lib/trino-plugin-toolkit/src/test/java/io/trino/plugin/base/logging/TestFormatInterpolator.java

@@ -0,0 +1,93 @@
+/*
+ * 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.base.logging;
+
+import org.testng.annotations.Test;
+
+import static org.assertj.core.api.AssertionsForInterfaceTypes.assertThat;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertTrue;
+
+public class TestFormatInterpolator
+{
+    @Test
+    public void testNullInterpolation()
+    {
+        FormatInterpolator<String> interpolator = new FormatInterpolator<>(null, SingleTestValue.values());
+        assertThat(interpolator.interpolate("!")).isEqualTo("");
+    }
+
+    @Test
+    public void testSingleValueInterpolation()
+    {
+        FormatInterpolator<String> interpolator = new FormatInterpolator<>("TEST_VALUE is $TEST_VALUE", SingleTestValue.values());
+        assertThat(interpolator.interpolate("!")).isEqualTo("TEST_VALUE is singleValue!");
+    }
+
+    @Test
+    public void testMultipleValueInterpolation()
+    {
+        FormatInterpolator<String> interpolator = new FormatInterpolator<>("TEST_VALUE is $TEST_VALUE and ANOTHER_VALUE is $ANOTHER_VALUE", MultipleTestValues.values());
+        assertThat(interpolator.interpolate("!")).isEqualTo("TEST_VALUE is first! and ANOTHER_VALUE is second!");
+    }
+
+    @Test
+    public void testUnknownValueInterpolation()
+    {
+        FormatInterpolator<String> interpolator = new FormatInterpolator<>("UNKNOWN_VALUE is $UNKNOWN_VALUE", MultipleTestValues.values());
+        assertThat(interpolator.interpolate("!")).isEqualTo("UNKNOWN_VALUE is $UNKNOWN_VALUE");
+    }
+
+    @Test
+    public void testValidation()
+    {
+        assertFalse(FormatInterpolator.hasValidPlaceholders("$UNKNOWN_VALUE", MultipleTestValues.values()));
+        assertTrue(FormatInterpolator.hasValidPlaceholders("$TEST_VALUE", MultipleTestValues.values()));
+        assertFalse(FormatInterpolator.hasValidPlaceholders("$TEST_VALUE and $UNKNOWN_VALUE", MultipleTestValues.values()));
+        assertTrue(FormatInterpolator.hasValidPlaceholders("$TEST_VALUE and $ANOTHER_VALUE", MultipleTestValues.values()));
+        assertTrue(FormatInterpolator.hasValidPlaceholders("$TEST_VALUE and $TEST_VALUE", MultipleTestValues.values()));
+    }
+
+    public enum SingleTestValue
+            implements FormatInterpolator.InterpolatedValue<String>
+    {
+        TEST_VALUE;
+
+        @Override
+        public String value(String value)
+        {
+            return "singleValue" + value;
+        }
+    }
+
+    public enum MultipleTestValues
+            implements FormatInterpolator.InterpolatedValue<String>
+    {
+        TEST_VALUE("first"),
+        ANOTHER_VALUE("second");
+
+        private final String value;
+
+        MultipleTestValues(String value)
+        {
+            this.value = value;
+        }
+
+        @Override
+        public String value(String value)
+        {
+            return this.value + value;
+        }
+    }
+}

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/logging/FormatBasedRemoteQueryModifier.java

@@ -14,93 +14,36 @@
 package io.trino.plugin.jdbc.logging;
 
 import com.google.inject.Inject;
+import io.trino.plugin.base.logging.FormatInterpolator;
+import io.trino.plugin.base.logging.SessionInterpolatedValues;
 import io.trino.spi.TrinoException;
 import io.trino.spi.connector.ConnectorSession;
 
-import java.util.function.Function;
-import java.util.function.Predicate;
-import java.util.regex.Pattern;
-
 import static com.google.common.base.Preconditions.checkState;
 import static io.trino.plugin.jdbc.JdbcErrorCode.JDBC_NON_TRANSIENT_ERROR;
-import static java.lang.String.format;
 import static java.util.Objects.requireNonNull;
 
 public class FormatBasedRemoteQueryModifier
         implements RemoteQueryModifier
 {
-    private final String commentFormat;
+    private final FormatInterpolator<ConnectorSession> interpolator;
 
     @Inject
     public FormatBasedRemoteQueryModifier(FormatBasedRemoteQueryModifierConfig config)
     {
-        this.commentFormat = requireNonNull(config, "config is null").getFormat();
+        String commentFormat = requireNonNull(config, "config is null").getFormat();
         checkState(!commentFormat.isBlank(), "comment format is blank");
+        this.interpolator = new FormatInterpolator<>(commentFormat, SessionInterpolatedValues.values());
     }
 
     @Override
     public String apply(ConnectorSession session, String query)
     {
-        String message = commentFormat;
-        for (PredefinedValue predefinedValue : PredefinedValue.values()) {
-            if (message.contains(predefinedValue.getPredefinedValueCode())) {
-                message = message.replaceAll(predefinedValue.getMatchCase(), predefinedValue.value(session));
-            }
-        }
-        return query + " /*" + message + "*/";
-    }
-
-    enum PredefinedValue
-    {
-        QUERY_ID(ConnectorSession::getQueryId),
-        SOURCE(new SanitizedValuesProvider(session -> session.getSource().orElse(""), "$SOURCE")),
-        USER(ConnectorSession::getUser),
-        TRACE_TOKEN(new SanitizedValuesProvider(session -> session.getTraceToken().orElse(""), "$TRACE_TOKEN"));
-
-        private final Function<ConnectorSession, String> valueProvider;
-
-        PredefinedValue(Function<ConnectorSession, String> valueProvider)
-        {
-            this.valueProvider = valueProvider;
+        try {
+            return query + " /*" + interpolator.interpolate(session) + "*/";
         }
-
-        String getMatchCase()
-        {
-            return "\\$" + this.name();
-        }
-
-        String getPredefinedValueCode()
-        {
-            return "$" + this.name();
-        }
-
-        String value(ConnectorSession session)
-        {
-            return valueProvider.apply(session);
-        }
-    }
-
-    private static class SanitizedValuesProvider
-            implements Function<ConnectorSession, String>
-    {
-        private static final Predicate<String> VALIDATION_MATCHER = Pattern.compile("^[\\w_-]*$").asMatchPredicate();
-        private final Function<ConnectorSession, String> valueProvider;
-        private final String name;
-
-        private SanitizedValuesProvider(Function<ConnectorSession, String> valueProvider, String name)
-        {
-            this.valueProvider = requireNonNull(valueProvider, "valueProvider is null");
-            this.name = requireNonNull(name, "name is null");
-        }
-
-        @Override
-        public String apply(ConnectorSession session)
-        {
-            String value = valueProvider.apply(session);
-            if (VALIDATION_MATCHER.test(value)) {
-                return value;
-            }
-            throw new TrinoException(JDBC_NON_TRANSIENT_ERROR, format("Passed value %s as %s does not meet security criteria. It can contain only letters, digits, underscores and hyphens", value, name));
+        catch (SecurityException e) {
+            throw new TrinoException(JDBC_NON_TRANSIENT_ERROR, e.getMessage());
         }
     }
 }

plugin/trino-base-jdbc/src/main/java/io/trino/plugin/jdbc/logging/FormatBasedRemoteQueryModifierConfig.java

@@ -15,22 +15,14 @@
 
 import io.airlift.configuration.Config;
 import io.airlift.configuration.ConfigDescription;
-import io.trino.plugin.jdbc.logging.FormatBasedRemoteQueryModifier.PredefinedValue;
+import io.trino.plugin.base.logging.SessionInterpolatedValues;
 
 import javax.validation.constraints.AssertTrue;
 
-import java.util.Arrays;
-import java.util.List;
-import java.util.regex.MatchResult;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
-
-import static java.util.stream.Collectors.joining;
+import static io.trino.plugin.base.logging.FormatInterpolator.hasValidPlaceholders;
 
 public class FormatBasedRemoteQueryModifierConfig
 {
-    private static final List<String> PREDEFINED_MATCHES = Arrays.stream(PredefinedValue.values()).map(PredefinedValue::getMatchCase).toList();
-    private static final Pattern VALIDATION_PATTERN = Pattern.compile("[\\w ,=]|" + String.join("|", PREDEFINED_MATCHES));
     private String format = "";
 
     @Config("query.comment-format")
@@ -49,10 +41,6 @@ public String getFormat()
     @AssertTrue(message = "Incorrect format it may consist of only letters, digits, underscores, commas, spaces, equal signs and predefined values")
     boolean isFormatValid()
     {
-        Matcher matcher = VALIDATION_PATTERN.matcher(format);
-        return matcher.results()
-                .map(MatchResult::group)
-                .collect(joining())
-                .equals(format);
+        return hasValidPlaceholders(format, SessionInterpolatedValues.values());
     }
 }