Skip to content

Commit

Permalink
Convert ClientSession to follow builder pattern
Browse files Browse the repository at this point in the history
  • Loading branch information
@akshayrai @phd3
akshayrai authored and phd3 committed Sep 7, 2022
1 parent df1dde1 commit 77f433a
Show file tree
Hide file tree
Showing 8 changed files with 199 additions and 170 deletions.
41 changes: 19 additions & 22 deletions client/trino-cli/src/main/java/io/trino/cli/ClientOptions.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,6 @@
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Strings.nullToEmpty;
import static io.trino.client.KerberosUtil.defaultCredentialCachePath;
import static java.util.Collections.emptyMap;
import static java.util.Locale.ENGLISH;
import static java.util.Objects.requireNonNull;
import static picocli.CommandLine.Option;
Expand Down Expand Up @@ -219,27 +218,25 @@ public String getKeyMap()

public ClientSession toClientSession()
{
return new ClientSession(
parseServer(server),
user,
sessionUser,
source,
Optional.ofNullable(traceToken),
parseClientTags(nullToEmpty(clientTags)),
clientInfo,
catalog,
schema,
null,
timeZone,
Locale.getDefault(),
toResourceEstimates(resourceEstimates),
toProperties(sessionProperties),
emptyMap(),
emptyMap(),
toExtraCredentials(extraCredentials),
null,
clientRequestTimeout,
disableCompression);
return ClientSession.builder()
.server(parseServer(server))
.principal(user)
.user(sessionUser)
.source(source)
.traceToken(Optional.ofNullable(traceToken))
.clientTags(parseClientTags(nullToEmpty(clientTags)))
.clientInfo(clientInfo)
.catalog(catalog)
.schema(schema)
.timeZone(timeZone)
.locale(Locale.getDefault())
.resourceEstimates(toResourceEstimates(resourceEstimates))
.properties(toProperties(sessionProperties))
.credentials(toExtraCredentials(extraCredentials))
.transactionId(null)
.clientRequestTimeout(clientRequestTimeout)
.compressionDisabled(disableCompression)
.build();
}

public static URI parseServer(String server)
Expand Down
14 changes: 7 additions & 7 deletions client/trino-cli/src/main/java/io/trino/cli/Console.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -366,8 +366,8 @@ private static boolean process(
// update catalog and schema if present
if (query.getSetCatalog().isPresent() || query.getSetSchema().isPresent()) {
session = ClientSession.builder(session)
.withCatalog(query.getSetCatalog().orElse(session.getCatalog()))
.withSchema(query.getSetSchema().orElse(session.getSchema()))
.catalog(query.getSetCatalog().orElse(session.getCatalog()))
.schema(query.getSetSchema().orElse(session.getSchema()))
.build();
}

Expand All @@ -379,35 +379,35 @@ private static boolean process(
ClientSession.Builder builder = ClientSession.builder(session);

if (query.getStartedTransactionId() != null) {
builder = builder.withTransactionId(query.getStartedTransactionId());
builder = builder.transactionId(query.getStartedTransactionId());
}

// update path if present
if (query.getSetPath().isPresent()) {
builder = builder.withPath(query.getSetPath().get());
builder = builder.path(query.getSetPath().get());
}

// update session properties if present
if (!query.getSetSessionProperties().isEmpty() || !query.getResetSessionProperties().isEmpty()) {
Map<String, String> sessionProperties = new HashMap<>(session.getProperties());
sessionProperties.putAll(query.getSetSessionProperties());
sessionProperties.keySet().removeAll(query.getResetSessionProperties());
builder = builder.withProperties(sessionProperties);
builder = builder.properties(sessionProperties);
}

// update session roles
if (!query.getSetRoles().isEmpty()) {
Map<String, ClientSelectedRole> roles = new HashMap<>(session.getRoles());
roles.putAll(query.getSetRoles());
builder = builder.withRoles(roles);
builder = builder.roles(roles);
}

// update prepared statements if present
if (!query.getAddedPreparedStatements().isEmpty() || !query.getDeallocatedPreparedStatements().isEmpty()) {
Map<String, String> preparedStatements = new HashMap<>(session.getPreparedStatements());
preparedStatements.putAll(query.getAddedPreparedStatements());
preparedStatements.keySet().removeAll(query.getDeallocatedPreparedStatements());
builder = builder.withPreparedStatements(preparedStatements);
builder = builder.preparedStatements(preparedStatements);
}

session = builder.build();
Expand Down
36 changes: 13 additions & 23 deletions client/trino-cli/src/test/java/io/trino/cli/TestQueryRunner.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,6 @@
package io.trino.cli;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import io.airlift.json.JsonCodec;
import io.airlift.units.Duration;
import io.trino.client.ClientSession;
Expand Down Expand Up @@ -102,27 +100,19 @@ public void testCookie()

static ClientSession createClientSession(MockWebServer server)
{
return new ClientSession(
server.url("/").uri(),
Optional.of("user"),
Optional.empty(),
"source",
Optional.empty(),
ImmutableSet.of(),
"clientInfo",
"catalog",
"schema",
null,
ZoneId.of("America/Los_Angeles"),
Locale.ENGLISH,
ImmutableMap.of(),
ImmutableMap.of(),
ImmutableMap.of(),
ImmutableMap.of(),
ImmutableMap.of(),
null,
new Duration(2, MINUTES),
true);
return ClientSession.builder()
.server(server.url("/").uri())
.principal(Optional.of("user"))
.source("source")
.clientInfo("clientInfo")
.catalog("catalog")
.schema("schema")
.timeZone(ZoneId.of("America/Los_Angeles"))
.locale(Locale.ENGLISH)
.transactionId(null)
.clientRequestTimeout(new Duration(2, MINUTES))
.compressionDisabled(true)
.build();
}

static String createResults(MockWebServer server)
Expand Down
129 changes: 98 additions & 31 deletions client/trino-client/src/main/java/io/trino/client/ClientSession.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,11 @@ public class ClientSession
private final Duration clientRequestTimeout;
private final boolean compressionDisabled;

public static Builder builder()
{
return new Builder();
}

public static Builder builder(ClientSession clientSession)
{
return new Builder(clientSession);
Expand All @@ -62,11 +67,11 @@ public static Builder builder(ClientSession clientSession)
public static ClientSession stripTransactionId(ClientSession session)
{
return ClientSession.builder(session)
.withoutTransactionId()
.transactionId(null)
.build();
}

public ClientSession(
private ClientSession(
URI server,
Optional<String> principal,
Optional<String> user,
Expand All @@ -89,8 +94,8 @@ public ClientSession(
boolean compressionDisabled)
{
this.server = requireNonNull(server, "server is null");
this.principal = principal;
this.user = user;
this.principal = requireNonNull(principal, "principal is null");
this.user = requireNonNull(user, "user is null");
this.source = source;
this.traceToken = requireNonNull(traceToken, "traceToken is null");
this.clientTags = ImmutableSet.copyOf(requireNonNull(clientTags, "clientTags is null"));
Expand Down Expand Up @@ -270,26 +275,28 @@ public String toString()
public static final class Builder
{
private URI server;
private Optional<String> principal;
private Optional<String> user;
private Optional<String> principal = Optional.empty();
private Optional<String> user = Optional.empty();
private String source;
private Optional<String> traceToken;
private Set<String> clientTags;
private Optional<String> traceToken = Optional.empty();
private Set<String> clientTags = ImmutableSet.of();
private String clientInfo;
private String catalog;
private String schema;
private String path;
private ZoneId timeZone;
private Locale locale;
private Map<String, String> resourceEstimates;
private Map<String, String> properties;
private Map<String, String> preparedStatements;
private Map<String, ClientSelectedRole> roles;
private Map<String, String> credentials;
private Map<String, String> resourceEstimates = ImmutableMap.of();
private Map<String, String> properties = ImmutableMap.of();
private Map<String, String> preparedStatements = ImmutableMap.of();
private Map<String, ClientSelectedRole> roles = ImmutableMap.of();
private Map<String, String> credentials = ImmutableMap.of();
private String transactionId;
private Duration clientRequestTimeout;
private boolean compressionDisabled;

private Builder() {}

private Builder(ClientSession clientSession)
{
requireNonNull(clientSession, "clientSession is null");
Expand All @@ -315,61 +322,121 @@ private Builder(ClientSession clientSession)
compressionDisabled = clientSession.isCompressionDisabled();
}

public Builder withCatalog(String catalog)
public Builder server(URI server)
{
this.server = server;
return this;
}

public Builder user(Optional<String> user)
{
this.user = user;
return this;
}

public Builder principal(Optional<String> principal)
{
this.principal = principal;
return this;
}

public Builder source(String source)
{
this.source = source;
return this;
}

public Builder traceToken(Optional<String> traceToken)
{
this.traceToken = traceToken;
return this;
}

public Builder clientTags(Set<String> clientTags)
{
this.clientTags = clientTags;
return this;
}

public Builder clientInfo(String clientInfo)
{
this.clientInfo = clientInfo;
return this;
}

public Builder catalog(String catalog)
{
this.catalog = catalog;
return this;
}

public Builder schema(String schema)
{
this.schema = schema;
return this;
}

public Builder path(String path)
{
this.path = path;
return this;
}

public Builder timeZone(ZoneId timeZone)
{
this.catalog = requireNonNull(catalog, "catalog is null");
this.timeZone = timeZone;
return this;
}

public Builder withSchema(String schema)
public Builder locale(Locale locale)
{
this.schema = requireNonNull(schema, "schema is null");
this.locale = locale;
return this;
}

public Builder withPath(String path)
public Builder resourceEstimates(Map<String, String> resourceEstimates)
{
this.path = requireNonNull(path, "path is null");
this.resourceEstimates = resourceEstimates;
return this;
}

public Builder withProperties(Map<String, String> properties)
public Builder properties(Map<String, String> properties)
{
this.properties = requireNonNull(properties, "properties is null");
this.properties = properties;
return this;
}

public Builder withRoles(Map<String, ClientSelectedRole> roles)
public Builder roles(Map<String, ClientSelectedRole> roles)
{
this.roles = roles;
return this;
}

public Builder withCredentials(Map<String, String> credentials)
public Builder credentials(Map<String, String> credentials)
{
this.credentials = requireNonNull(credentials, "credentials is null");
this.credentials = credentials;
return this;
}

public Builder withPreparedStatements(Map<String, String> preparedStatements)
public Builder preparedStatements(Map<String, String> preparedStatements)
{
this.preparedStatements = requireNonNull(preparedStatements, "preparedStatements is null");
this.preparedStatements = preparedStatements;
return this;
}

public Builder withTransactionId(String transactionId)
public Builder transactionId(String transactionId)
{
this.transactionId = requireNonNull(transactionId, "transactionId is null");
this.transactionId = transactionId;
return this;
}

public Builder withoutTransactionId()
public Builder clientRequestTimeout(Duration clientRequestTimeout)
{
this.transactionId = null;
this.clientRequestTimeout = clientRequestTimeout;
return this;
}

public Builder withCompressionDisabled(boolean compressionDisabled)
public Builder compressionDisabled(boolean compressionDisabled)
{
this.compressionDisabled = compressionDisabled;
return this;
Expand Down
Loading

0 comments on commit 77f433a

Please sign in to comment.