Skip to content

feat: Preserve selectedUser (identity) in Write Queries (#27360) #27360

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
shelton408 merged 1 commit into from
Mar 20, 2026
Merged

feat: Preserve selectedUser (identity) in Write Queries (#27360) #27360

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion presto-main-base/src/main/java/com/facebook/presto/Session.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -591,7 +591,9 @@ public SessionRepresentation toSessionRepresentation()
unprocessedCatalogProperties,
identity.getRoles(),
preparedStatements,
sessionFunctions);
sessionFunctions,
identity.getSelectedUser(),
identity.getReasonForSelect());
}

@Override
Expand Down
26 changes: 23 additions & 3 deletions presto-main-base/src/main/java/com/facebook/presto/SessionRepresentation.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,8 @@ public final class SessionRepresentation
private final Map<String, SelectedRole> roles;
private final Map<String, String> preparedStatements;
private final Map<SqlFunctionId, SqlInvokedFunction> sessionFunctions;
private final Optional<String> selectedUser;
private final Optional<String> reasonForSelect;

@ThriftConstructor
@JsonCreator
Expand All @@ -95,7 +97,9 @@ public SessionRepresentation(
@JsonProperty("unprocessedCatalogProperties") Map<String, Map<String, String>> unprocessedCatalogProperties,
@JsonProperty("roles") Map<String, SelectedRole> roles,
@JsonProperty("preparedStatements") Map<String, String> preparedStatements,
@JsonProperty("sessionFunctions") Map<SqlFunctionId, SqlInvokedFunction> sessionFunctions)
@JsonProperty("sessionFunctions") Map<SqlFunctionId, SqlInvokedFunction> sessionFunctions,
@JsonProperty("selectedUser") Optional<String> selectedUser,
@JsonProperty("reasonForSelect") Optional<String> reasonForSelect)
{
this.queryId = requireNonNull(queryId, "queryId is null");
this.transactionId = requireNonNull(transactionId, "transactionId is null");
Expand All @@ -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<ConnectorId, Map<String, String>> catalogPropertiesBuilder = ImmutableMap.builder();
for (Entry<ConnectorId, Map<String, String>> entry : catalogProperties.entrySet()) {
Expand Down Expand Up @@ -293,6 +299,20 @@ public Map<SqlFunctionId, SqlInvokedFunction> getSessionFunctions()
return sessionFunctions;
}

@ThriftField(24)
@JsonProperty
public Optional<String> getSelectedUser()
{
return selectedUser;
}

@ThriftField(25)
@JsonProperty
public Optional<String> getReasonForSelect()
{
return reasonForSelect;
}

public Session toSession(SessionPropertyManager sessionPropertyManager)
{
return toSession(sessionPropertyManager, emptyMap(), emptyMap());
Expand All @@ -315,8 +335,8 @@ public Session toSession(SessionPropertyManager sessionPropertyManager, Map<Stri
roles,
extraCredentials,
extraAuthenticators,
Optional.empty(),
Optional.empty()),
selectedUser,
reasonForSelect),
source,
catalog,
schema,
Expand Down
4 changes: 4 additions & 0 deletions presto-native-execution/presto_cpp/main/thrift/ProtocolToThrift.cpp
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

2 changes: 2 additions & 0 deletions presto-native-execution/presto_cpp/main/thrift/presto_thrift.thrift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -527,6 +527,8 @@ struct SessionRepresentation {
21: map<string, SelectedRole> roles;
22: map<string, string> preparedStatements;
23: map<SqlFunctionId, SqlInvokedFunction> sessionFunctions;
24: optional string selectedUser;
25: optional string reasonForSelect;
}
struct SelectedRole {
1: Type type;
Expand Down
28 changes: 28 additions & 0 deletions presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2144,6 +2144,20 @@ void to_json(json& j, const SessionRepresentation& p) {
"SessionRepresentation",
"Map<SqlFunctionId, SqlInvokedFunction>",
"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) {
Expand Down Expand Up @@ -2277,6 +2291,20 @@ void from_json(const json& j, SessionRepresentation& p) {
"SessionRepresentation",
"Map<SqlFunctionId, SqlInvokedFunction>",
"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 {
Expand Down
2 changes: 2 additions & 0 deletions presto-native-execution/presto_cpp/presto_protocol/core/presto_protocol_core.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -706,6 +706,8 @@ struct SessionRepresentation {
Map<String, SelectedRole> roles = {};
Map<String, String> preparedStatements = {};
Map<SqlFunctionId, SqlInvokedFunction> sessionFunctions = {};
std::shared_ptr<String> selectedUser = {};
std::shared_ptr<String> reasonForSelect = {};
};
void to_json(json& j, const SessionRepresentation& p);
void from_json(const json& j, SessionRepresentation& p);
Expand Down
4 changes: 2 additions & 2 deletions ...rc/main/java/com/facebook/presto/spark/execution/task/PrestoSparkTaskExecutorFactory.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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(),
Expand Down
Loading