From 8b77f99ff394bc02f6cc994c51105dcadffa81ef Mon Sep 17 00:00:00 2001 From: Shelton Cai Date: Wed, 18 Mar 2026 17:03:43 -0700 Subject: [PATCH] Preserve Identity in Session Representation (#27360) Summary: Pull Request resolved: https://github.com/prestodb/presto/pull/27360 WARNING: i ignored structured annotations using thrift2json.py for presto-native in this PR. Thrift generation was having trouble reading certain annotations without this, but I'm not sure if this is will cause any issues in protocol generation. If anyone has context please let me know. When plumbing identity to metastore for write/coordinator related queries (e.g. Insert/Delete), we convert Session to a SessionRepresentation for serialization, however this representation uses Optional.Empty as placeholders for identity. The identity is lost in the following steps 1.session is created with user identity 2. serialize session representation, drops identity in this serialization 3. deserialize the session representation and use this as the session for the query 4. uses deserialized session representation in metastore This poses an issue in cases where we want to pass the selected identity to an external metastore service for authentication. Changes: Update Session/Session Representation to include selectedUser and reasonForSelect in the conversion methods. Include both as thrift fields in SessionRepresentation CPP protocol and thrift updates to match the change to session representation thrift fields. thrift2json.py change to ignore certain annotations *** not sure if this should be committed. Included the fields in conversion back to Session for Spark as well for consistency. Differential Revision: D97002318 --- .../java/com/facebook/presto/Session.java | 4 ++- .../presto/SessionRepresentation.java | 26 +++++++++++++++-- .../main/thrift/ProtocolToThrift.cpp | 4 +++ .../main/thrift/presto_thrift.thrift | 2 ++ .../core/presto_protocol_core.cpp | 28 +++++++++++++++++++ .../core/presto_protocol_core.h | 2 ++ .../task/PrestoSparkTaskExecutorFactory.java | 4 +-- 7 files changed, 64 insertions(+), 6 deletions(-) diff --git a/presto-main-base/src/main/java/com/facebook/presto/Session.java b/presto-main-base/src/main/java/com/facebook/presto/Session.java index 63219c39082e9..29b1b247ce7b0 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/Session.java +++ b/presto-main-base/src/main/java/com/facebook/presto/Session.java @@ -591,7 +591,9 @@ public SessionRepresentation toSessionRepresentation() unprocessedCatalogProperties, identity.getRoles(), preparedStatements, - sessionFunctions); + sessionFunctions, + identity.getSelectedUser(), + identity.getReasonForSelect()); } @Override diff --git a/presto-main-base/src/main/java/com/facebook/presto/SessionRepresentation.java b/presto-main-base/src/main/java/com/facebook/presto/SessionRepresentation.java index f8203723ec44b..d4b15ec32d043 100644 --- a/presto-main-base/src/main/java/com/facebook/presto/SessionRepresentation.java +++ b/presto-main-base/src/main/java/com/facebook/presto/SessionRepresentation.java @@ -69,6 +69,8 @@ public final class SessionRepresentation private final Map roles; private final Map preparedStatements; private final Map sessionFunctions; + private final Optional selectedUser; + private final Optional reasonForSelect; @ThriftConstructor @JsonCreator @@ -95,7 +97,9 @@ public SessionRepresentation( @JsonProperty("unprocessedCatalogProperties") Map> unprocessedCatalogProperties, @JsonProperty("roles") Map roles, @JsonProperty("preparedStatements") Map preparedStatements, - @JsonProperty("sessionFunctions") Map sessionFunctions) + @JsonProperty("sessionFunctions") Map sessionFunctions, + @JsonProperty("selectedUser") Optional selectedUser, + @JsonProperty("reasonForSelect") Optional reasonForSelect) { this.queryId = requireNonNull(queryId, "queryId is null"); this.transactionId = requireNonNull(transactionId, "transactionId is null"); @@ -118,6 +122,8 @@ public SessionRepresentation( this.roles = ImmutableMap.copyOf(roles); this.preparedStatements = ImmutableMap.copyOf(preparedStatements); this.sessionFunctions = ImmutableMap.copyOf(sessionFunctions); + this.selectedUser = selectedUser == null ? Optional.empty() : selectedUser; + this.reasonForSelect = reasonForSelect == null ? Optional.empty() : reasonForSelect; ImmutableMap.Builder> catalogPropertiesBuilder = ImmutableMap.builder(); for (Entry> entry : catalogProperties.entrySet()) { @@ -293,6 +299,20 @@ public Map getSessionFunctions() return sessionFunctions; } + @ThriftField(24) + @JsonProperty + public Optional getSelectedUser() + { + return selectedUser; + } + + @ThriftField(25) + @JsonProperty + public Optional getReasonForSelect() + { + return reasonForSelect; + } + public Session toSession(SessionPropertyManager sessionPropertyManager) { return toSession(sessionPropertyManager, emptyMap(), emptyMap()); @@ -315,8 +335,8 @@ public Session toSession(SessionPropertyManager sessionPropertyManager, Map roles; 22: map preparedStatements; 23: map sessionFunctions; + 24: optional string selectedUser; + 25: optional string reasonForSelect; } struct SelectedRole { 1: Type type; diff --git a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp index 39b1686cf0692..809f5ee4a2267 100644 --- a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp +++ b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp @@ -2144,6 +2144,20 @@ void to_json(json& j, const SessionRepresentation& p) { "SessionRepresentation", "Map", "sessionFunctions"); + to_json_key( + j, + "selectedUser", + p.selectedUser, + "SessionRepresentation", + "String", + "selectedUser"); + to_json_key( + j, + "reasonForSelect", + p.reasonForSelect, + "SessionRepresentation", + "String", + "reasonForSelect"); } void from_json(const json& j, SessionRepresentation& p) { @@ -2277,6 +2291,20 @@ void from_json(const json& j, SessionRepresentation& p) { "SessionRepresentation", "Map", "sessionFunctions"); + from_json_key( + j, + "selectedUser", + p.selectedUser, + "SessionRepresentation", + "String", + "selectedUser"); + from_json_key( + j, + "reasonForSelect", + p.reasonForSelect, + "SessionRepresentation", + "String", + "reasonForSelect"); } } // namespace facebook::presto::protocol namespace facebook::presto::protocol { diff --git a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h index a20589ae571ff..f5fa9ed15236b 100644 --- a/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h +++ b/presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h @@ -706,6 +706,8 @@ struct SessionRepresentation { Map roles = {}; Map preparedStatements = {}; Map sessionFunctions = {}; + std::shared_ptr selectedUser = {}; + std::shared_ptr reasonForSelect = {}; }; void to_json(json& j, const SessionRepresentation& p); void from_json(const json& j, SessionRepresentation& p); diff --git a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/task/PrestoSparkTaskExecutorFactory.java b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/task/PrestoSparkTaskExecutorFactory.java index 0123bf3a8bcd2..fffdce9a1bdf9 100644 --- a/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/task/PrestoSparkTaskExecutorFactory.java +++ b/presto-spark-base/src/main/java/com/facebook/presto/spark/execution/task/PrestoSparkTaskExecutorFactory.java @@ -661,8 +661,8 @@ private Session createSessionWithExtraSessionProperties( sessionRepresentation.getRoles(), extraCredentials, extraAuthenticators, - Optional.empty(), - Optional.empty()), + sessionRepresentation.getSelectedUser(), + sessionRepresentation.getReasonForSelect()), sessionRepresentation.getSource(), sessionRepresentation.getCatalog(), sessionRepresentation.getSchema(),