Migrate query state info resource to re2j regex and add test

unix280 · Dec 17, 2024 · b2da48d · b2da48d
1 parent 50a7425
commit b2da48d
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 4 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/server/QueryStateInfoResource.java b/presto-main/src/main/java/com/facebook/presto/server/QueryStateInfoResource.java
@@ -20,6 +20,8 @@
 import com.facebook.presto.resourcemanager.ResourceManagerProxy;
 import com.facebook.presto.spi.QueryId;
 import com.facebook.presto.spi.resourceGroups.ResourceGroupId;
+import com.google.re2j.Pattern;
+import io.airlift.slice.Slices;
 
 import javax.annotation.security.RolesAllowed;
 import javax.inject.Inject;
@@ -45,7 +47,6 @@
 import java.util.NoSuchElementException;
 import java.util.Optional;
 import java.util.OptionalInt;
-import java.util.regex.Pattern;
 
 import static com.facebook.presto.execution.QueryState.QUEUED;
 import static com.facebook.presto.server.QueryStateInfo.createQueryStateInfo;
@@ -107,7 +108,7 @@ public void getQueryStateInfos(
 
             List<QueryStateInfo> queryStateInfos = queryInfos.stream()
                     .filter(queryInfo -> includeAllQueries || !queryInfo.getState().isDone())
-                    .filter(queryInfo -> userPattern.map(pattern -> pattern.matcher(queryInfo.getSession().getUser()).matches()).orElse(true))
+                    .filter(queryInfo -> userPattern.map(pattern -> pattern.matcher(Slices.utf8Slice(queryInfo.getSession().getUser())).matches()).orElse(true))
                     .map(queryInfo -> getQueryStateInfo(
                             queryInfo,
                             includeAllQueryProgressStats,

diff --git a/presto-main/src/test/java/com/facebook/presto/server/TestQueryStateInfoResource.java b/presto-main/src/test/java/com/facebook/presto/server/TestQueryStateInfoResource.java
@@ -40,6 +40,7 @@
 import static com.facebook.presto.testing.assertions.Assert.assertEquals;
 import static java.nio.charset.StandardCharsets.UTF_8;
 import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertThrows;
 import static org.testng.Assert.assertTrue;
 
 @Test(singleThreaded = true)
@@ -80,12 +81,20 @@ public void setup()
         QueryResults queryResults2 = client.execute(request2, createJsonResponseHandler(jsonCodec(QueryResults.class)));
         client.execute(prepareGet().setUri(queryResults2.getNextUri()).build(), createJsonResponseHandler(QUERY_RESULTS_JSON_CODEC));
 
+        Request request3 = preparePost()
+                .setUri(uriBuilderFrom(server.getBaseUrl()).replacePath("/v1/statement").build())
+                .setBodyGenerator(createStaticBodyGenerator(LONG_LASTING_QUERY, UTF_8))
+                .setHeader(PRESTO_USER, "alt3")
+                .build();
+        QueryResults queryResults3 = client.execute(request3, createJsonResponseHandler(jsonCodec(QueryResults.class)));
+        client.execute(prepareGet().setUri(queryResults3.getNextUri()).build(), createJsonResponseHandler(QUERY_RESULTS_JSON_CODEC));
+
         // queries are started in the background, so they may not all be immediately visible
         while (true) {
             List<BasicQueryInfo> queryInfos = client.execute(
                     prepareGet().setUri(uriBuilderFrom(server.getBaseUrl()).replacePath("/v1/query").build()).build(),
                     createJsonResponseHandler(listJsonCodec(BasicQueryInfo.class)));
-            if ((queryInfos.size() == 2) && queryInfos.stream().allMatch(info -> info.getState() == RUNNING)) {
+            if ((queryInfos.size() == 3) && queryInfos.stream().allMatch(info -> info.getState() == RUNNING)) {
                 break;
             }
         }
@@ -105,7 +114,7 @@ public void testGetAllQueryStateInfos()
                 prepareGet().setUri(server.resolve("/v1/queryState")).build(),
                 createJsonResponseHandler(listJsonCodec(QueryStateInfo.class)));
 
-        assertEquals(infos.size(), 2);
+        assertEquals(infos.size(), 3);
     }
 
     @Test
@@ -138,6 +147,37 @@ public void testGetQueryStateInfo()
         assertNotNull(info);
     }
 
+    @Test
+    public void testQueryStateInfoRegexMatching()
+    {
+        // Match strings containing "user"
+        List<QueryStateInfo> infos = client.execute(
+                prepareGet()
+                        .setUri(server.resolve("/v1/queryState?user=user%5Cd")) // Encoded user\d
+                        .build(),
+                createJsonResponseHandler(listJsonCodec(QueryStateInfo.class)));
+        assertEquals(infos.size(), 2);
+
+        // Match strings ending in a digit
+        infos = client.execute(
+                prepareGet()
+                        .setUri(server.resolve("/v1/queryState?user=%2E%2A%5Cd")) // Encoded .*\d
+                        .build(),
+                createJsonResponseHandler(listJsonCodec(QueryStateInfo.class)));
+        assertEquals(infos.size(), 3);
+
+        // Non supported version of previous pattern
+        infos = client.execute(
+                prepareGet()
+                        .setUri(server.resolve("/v1/queryState?user=%2E%7B%7D%5Cd")) // Encoded .{}\d
+                        .build(),
+                createJsonResponseHandler(listJsonCodec(QueryStateInfo.class)));
+        assertEquals(infos.size(), 0);
+
+        Request failingReq = prepareGet().setUri(server.resolve("/v1/queryState?user=%2A%5Cd")).build(); // Encoded *\d
+        assertThrows(() -> client.execute(failingReq, createJsonResponseHandler(listJsonCodec(QueryStateInfo.class))));
+    }
+
     @Test(expectedExceptions = {UnexpectedResponseException.class}, expectedExceptionsMessageRegExp = "Expected response code .*, but was 404")
     public void testGetQueryStateInfoNo()
     {