Include enabled roles in QueryContext for event listeners

It's odd that it included the principal and groups, but not enabled roles.
trinodb · Sep 15, 2023 · 7d36ed5 · 7d36ed5
1 parent 0c6fac4
commit 7d36ed5
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 0 deletions.
diff --git a/core/trino-main/src/main/java/io/trino/event/QueryMonitor.java b/core/trino-main/src/main/java/io/trino/event/QueryMonitor.java
@@ -345,6 +345,7 @@ private QueryContext createQueryContext(SessionRepresentation session, Optional<
                 session.getUser(),
                 session.getOriginalUser(),
                 session.getPrincipal(),
+                session.getEnabledRoles(),
                 session.getGroups(),
                 session.getTraceToken(),
                 session.getRemoteUserAddress(),

diff --git a/core/trino-spi/src/main/java/io/trino/spi/eventlistener/QueryContext.java b/core/trino-spi/src/main/java/io/trino/spi/eventlistener/QueryContext.java
@@ -34,6 +34,7 @@ public class QueryContext
     private final String user;
     private final String originalUser;
     private final Optional<String> principal;
+    private final Set<String> enabledRoles;
     private final Set<String> groups;
     private final Optional<String> traceToken;
     private final Optional<String> remoteClientAddress;
@@ -65,6 +66,7 @@ public QueryContext(
             String user,
             String originalUser,
             Optional<String> principal,
+            Set<String> enabledRoles,
             Set<String> groups,
             Optional<String> traceToken,
             Optional<String> remoteClientAddress,
@@ -87,6 +89,7 @@ public QueryContext(
         this.user = requireNonNull(user, "user is null");
         this.originalUser = requireNonNull(originalUser, "originalUser is null");
         this.principal = requireNonNull(principal, "principal is null");
+        this.enabledRoles = requireNonNull(enabledRoles, "enabledRoles is null");
         this.groups = requireNonNull(groups, "groups is null");
         this.traceToken = requireNonNull(traceToken, "traceToken is null");
         this.remoteClientAddress = requireNonNull(remoteClientAddress, "remoteClientAddress is null");
@@ -125,6 +128,12 @@ public Optional<String> getPrincipal()
         return principal;
     }
 
+    @JsonProperty
+    public Set<String> getEnabledRoles()
+    {
+        return enabledRoles;
+    }
+
     @JsonProperty
     public Set<String> getGroups()
     {

diff --git a/...no-http-event-listener/src/test/java/io/trino/plugin/httpquery/TestHttpEventListener.java b/...no-http-event-listener/src/test/java/io/trino/plugin/httpquery/TestHttpEventListener.java
@@ -107,6 +107,7 @@ public class TestHttpEventListener
                 "user",
                 "originalUser",
                 Optional.of("principal"),
+                Set.of(), // enabledRoles
                 Set.of(), // groups
                 Optional.empty(), // traceToken
                 Optional.empty(), // remoteClientAddress

diff --git a/...nt-listener/src/test/java/io/trino/plugin/eventlistener/mysql/TestMysqlEventListener.java b/...nt-listener/src/test/java/io/trino/plugin/eventlistener/mysql/TestMysqlEventListener.java
@@ -134,6 +134,7 @@ public class TestMysqlEventListener
             "user",
             "originalUser",
             Optional.of("principal"),
+            Set.of("role1", "role2"),
             Set.of("group1", "group2"),
             Optional.of("traceToken"),
             Optional.of("remoteAddress"),
@@ -288,6 +289,7 @@ public class TestMysqlEventListener
             "originalUser",
             Optional.empty(),
             Set.of(),
+            Set.of(),
             Optional.empty(),
             Optional.empty(),
             Optional.empty(),