From d588abed44e4c9b6940a5503335068b50de03eaf Mon Sep 17 00:00:00 2001
From: Krzysztof Sobolewski <krzysztof.sobolewski@starburstdata.com>
Date: Mon, 11 Sep 2023 11:40:22 +0200
Subject: [PATCH] Include enabled roles in QueryContext for event listeners

It's odd that it included the principal and groups, but not enabled
roles.
---
 .../src/main/java/io/trino/event/QueryMonitor.java       | 1 +
 .../java/io/trino/spi/eventlistener/QueryContext.java    | 9 +++++++++
 .../io/trino/plugin/httpquery/TestHttpEventListener.java | 1 +
 .../eventlistener/mysql/TestMysqlEventListener.java      | 2 ++
 4 files changed, 13 insertions(+)

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
index 7951fc9d80f3b..6f00c2baea713 100644
--- 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
index 6dc00edac896a..b0990299274b7 100644
--- 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/plugin/trino-http-event-listener/src/test/java/io/trino/plugin/httpquery/TestHttpEventListener.java b/plugin/trino-http-event-listener/src/test/java/io/trino/plugin/httpquery/TestHttpEventListener.java
index 5df2effd266ed..839dd73fa1b71 100644
--- a/plugin/trino-http-event-listener/src/test/java/io/trino/plugin/httpquery/TestHttpEventListener.java
+++ b/plugin/trino-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/plugin/trino-mysql-event-listener/src/test/java/io/trino/plugin/eventlistener/mysql/TestMysqlEventListener.java b/plugin/trino-mysql-event-listener/src/test/java/io/trino/plugin/eventlistener/mysql/TestMysqlEventListener.java
index 8f3dfbd9450ea..a30aaf4e525f9 100644
--- a/plugin/trino-mysql-event-listener/src/test/java/io/trino/plugin/eventlistener/mysql/TestMysqlEventListener.java
+++ b/plugin/trino-mysql-event-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(),