Add session's source and client info to QueryStateInfo

prestodb · Sep 12, 2017 · 9e88cc3 · 9e88cc3
1 parent 14453b8
commit 9e88cc3
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 0 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/server/QueryStateInfo.java b/presto-main/src/main/java/com/facebook/presto/server/QueryStateInfo.java
@@ -40,6 +40,8 @@ public class QueryStateInfo
     private final String query;
     private final DateTime createTime;
     private final String user;
+    private final Optional<String> source;
+    private final Optional<String> clientInfo;
     private final Optional<String> catalog;
     private final Optional<String> schema;
     private final Optional<List<ResourceGroupInfo>> resourceGroupChain;
@@ -53,6 +55,8 @@ public QueryStateInfo(
             @JsonProperty("query") String query,
             @JsonProperty("createTime") DateTime createTime,
             @JsonProperty("user") String user,
+            @JsonProperty("source") Optional<String> source,
+            @JsonProperty("clientInfo") Optional<String> clientInfo,
             @JsonProperty("catalog") Optional<String> catalog,
             @JsonProperty("schema") Optional<String> schema,
             @JsonProperty("resourceGroupChainInfo") Optional<List<ResourceGroupInfo>> resourceGroupChain,
@@ -64,6 +68,8 @@ public QueryStateInfo(
         this.query = requireNonNull(query, "query text is null");
         this.createTime = requireNonNull(createTime, "createTime is null");
         this.user = requireNonNull(user, "user is null");
+        this.source = requireNonNull(source, "source is null");
+        this.clientInfo = requireNonNull(clientInfo, "clientInfo is null");
         this.catalog = requireNonNull(catalog, "catalog is null");
         this.schema = requireNonNull(schema, "schema is null");
         requireNonNull(resourceGroupChain, "resourceGroupChain is null");
@@ -108,6 +114,8 @@ else if (queryInfo.getState() == RUNNING) {
                 queryInfo.getQuery(),
                 queryInfo.getQueryStats().getCreateTime(),
                 queryInfo.getSession().getUser(),
+                queryInfo.getSession().getSource(),
+                queryInfo.getSession().getClientInfo(),
                 queryInfo.getSession().getCatalog(),
                 queryInfo.getSession().getSchema(),
                 resourceGroups,
@@ -144,6 +152,18 @@ public String getUser()
         return user;
     }
 
+    @JsonProperty
+    public Optional<String> getSource()
+    {
+        return source;
+    }
+
+    @JsonProperty
+    public Optional<String> getClientInfo()
+    {
+        return clientInfo;
+    }
+
     @JsonProperty
     public Optional<String> getCatalog()
     {

diff --git a/presto-main/src/test/java/com/facebook/presto/server/TestQueryStateInfo.java b/presto-main/src/test/java/com/facebook/presto/server/TestQueryStateInfo.java
@@ -40,6 +40,7 @@
 import static com.facebook.presto.server.QueryStateInfo.createQueryStateInfo;
 import static com.facebook.presto.spi.resourceGroups.ResourceGroupState.CAN_QUEUE;
 import static com.facebook.presto.spi.resourceGroups.ResourceGroupState.CAN_RUN;
+import static com.facebook.presto.tpch.TpchMetadata.TINY_SCHEMA_NAME;
 import static io.airlift.units.DataSize.Unit.BYTE;
 import static java.util.concurrent.TimeUnit.DAYS;
 import static java.util.concurrent.TimeUnit.HOURS;
@@ -163,6 +164,8 @@ public void testQueryStateInfo()
                 createQueryInfo("query_root_b", QUEUED, "SELECT count(*) FROM t"),
                 Optional.of(groupRootB),
                 Optional.of(rootInfo));
+        assertEquals(infoForQueryQueuedOnRootB.getCatalog().get(), "tpch");
+        assertEquals(infoForQueryQueuedOnRootB.getSchema().get(), TINY_SCHEMA_NAME);
         assertEquals(infoForQueryQueuedOnRootB.getQuery(), "SELECT count(*) FROM t");
         assertEquals(infoForQueryQueuedOnRootB.getQueryId().toString(), "query_root_b");
         assertEquals(infoForQueryQueuedOnRootB.getQueryState(), QUEUED);

diff --git a/presto-main/src/test/java/com/facebook/presto/server/TestResourceGroupStateInfo.java b/presto-main/src/test/java/com/facebook/presto/server/TestResourceGroupStateInfo.java
@@ -56,6 +56,8 @@ public void testJsonRoundTrip()
                         "SELECT * FROM t",
                         DateTime.parse("2017-06-12T21:39:48.658Z"),
                         "test_user",
+                        Optional.of("source"),
+                        Optional.of("client_info"),
                         Optional.of("catalog"),
                         Optional.of("schema"),
                         Optional.empty(),
@@ -104,6 +106,8 @@ public void testJsonRoundTrip()
         assertEquals(queryStateInfo.getQuery(), "SELECT * FROM t");
         assertEquals(queryStateInfo.getCreateTime(), DateTime.parse("2017-06-12T21:39:48.658Z"));
         assertEquals(queryStateInfo.getUser(), "test_user");
+        assertEquals(queryStateInfo.getSource(), Optional.of("source"));
+        assertEquals(queryStateInfo.getClientInfo(), Optional.of("client_info"));
         assertEquals(queryStateInfo.getCatalog(), Optional.of("catalog"));
         assertEquals(queryStateInfo.getSchema(), Optional.of("schema"));
         assertEquals(queryStateInfo.getResourceGroupChain(), Optional.empty());