Skip to content

Commit 33032e7

Browse files
daindain
committed
Add source, remote client ip and user agent to session
1 parent bd6f2da commit 33032e7

File tree

24 files changed

+197
-34
lines changed

24 files changed

+197
-34
lines changed

‎presto-cli/src/main/java/com/facebook/presto/cli/ClientOptions.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,6 @@ public enum OutputFormat
4242

4343
public ClientSession toClientSession()
4444
{
45-
return new ClientSession(server, user, catalog, schema, debug);
45+
return new ClientSession(server, user, "presto-cli", catalog, schema, debug);
4646
}
4747
}

‎presto-client/src/main/java/com/facebook/presto/client/ClientSession.java‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,16 @@ public class ClientSession
1313
{
1414
private URI server;
1515
private String user;
16+
private String source;
1617
private String catalog;
1718
private String schema;
1819
private boolean debug;
1920

20-
public ClientSession(URI server, String user, String catalog, String schema, boolean debug)
21+
public ClientSession(URI server, String user, String source, String catalog, String schema, boolean debug)
2122
{
2223
this.server = checkNotNull(server, "server is null");
2324
this.user = user;
25+
this.source = source;
2426
this.catalog = catalog;
2527
this.schema = schema;
2628
this.debug = debug;
@@ -46,6 +48,16 @@ public void setUser(String user)
4648
this.user = user;
4749
}
4850

51+
public String getSource()
52+
{
53+
return source;
54+
}
55+
56+
public void setSource(String source)
57+
{
58+
this.source = source;
59+
}
60+
4961
public String getCatalog()
5062
{
5163
return catalog;

‎presto-client/src/main/java/com/facebook/presto/client/PrestoHeaders.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
public final class PrestoHeaders
77
{
88
public static final String PRESTO_USER = "X-Presto-User";
9+
public static final String PRESTO_SOURCE = "X-Presto-Source";
910
public static final String PRESTO_CATALOG = "X-Presto-Catalog";
1011
public static final String PRESTO_SCHEMA = "X-Presto-Schema";
1112

‎presto-client/src/main/java/com/facebook/presto/client/StatementClient.java‎

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,21 +1,22 @@
11
package com.facebook.presto.client;
22

33
import com.google.common.base.Charsets;
4+
import com.google.common.base.Objects;
45
import io.airlift.http.client.AsyncHttpClient;
56
import io.airlift.http.client.FullJsonResponseHandler;
67
import io.airlift.http.client.HttpStatus;
78
import io.airlift.http.client.Request;
89
import io.airlift.json.JsonCodec;
910

1011
import javax.annotation.concurrent.ThreadSafe;
11-
1212
import java.io.Closeable;
1313
import java.net.URI;
1414
import java.util.concurrent.atomic.AtomicBoolean;
1515
import java.util.concurrent.atomic.AtomicReference;
1616

1717
import static com.google.common.base.Preconditions.checkNotNull;
1818
import static com.google.common.base.Preconditions.checkState;
19+
import static com.google.common.net.HttpHeaders.USER_AGENT;
1920
import static com.google.common.util.concurrent.Uninterruptibles.sleepUninterruptibly;
2021
import static io.airlift.http.client.FullJsonResponseHandler.JsonResponse;
2122
import static io.airlift.http.client.FullJsonResponseHandler.createFullJsonResponseHandler;
@@ -36,6 +37,10 @@
3637
public class StatementClient
3738
implements Closeable
3839
{
40+
private static final String USER_AGENT_VALUE = StatementClient.class.getSimpleName() +
41+
"/" +
42+
Objects.firstNonNull(StatementClient.class.getPackage().getImplementationVersion(), "unknown");
43+
3944
private final AsyncHttpClient httpClient;
4045
private final FullJsonResponseHandler<QueryResults> responseHandler;
4146
private final boolean debug;
@@ -70,12 +75,16 @@ private static Request buildQueryRequest(ClientSession session, String query)
7075
if (session.getUser() != null) {
7176
builder.setHeader(PrestoHeaders.PRESTO_USER, session.getUser());
7277
}
78+
if (session.getSource() != null) {
79+
builder.setHeader(PrestoHeaders.PRESTO_SOURCE, session.getSource());
80+
}
7381
if (session.getCatalog() != null) {
7482
builder.setHeader(PrestoHeaders.PRESTO_CATALOG, session.getCatalog());
7583
}
7684
if (session.getSchema() != null) {
7785
builder.setHeader(PrestoHeaders.PRESTO_SCHEMA, session.getSchema());
7886
}
87+
builder.setHeader(USER_AGENT, USER_AGENT_VALUE);
7988

8089
return builder.build();
8190
}
@@ -129,7 +138,10 @@ public boolean advance()
129138
return false;
130139
}
131140

132-
Request request = prepareGet().setUri(current().getNextUri()).build();
141+
Request request = prepareGet()
142+
.setHeader(USER_AGENT, USER_AGENT_VALUE)
143+
.setUri(current().getNextUri())
144+
.build();
133145

134146
Exception cause = null;
135147
long start = System.nanoTime();
@@ -178,7 +190,10 @@ public boolean cancelLeafStage()
178190
return false;
179191
}
180192

181-
Request request = prepareDelete().setUri(uri).build();
193+
Request request = prepareDelete()
194+
.setHeader(USER_AGENT, USER_AGENT_VALUE)
195+
.setUri(uri)
196+
.build();
182197
StatusResponse status = httpClient.execute(request, createStatusResponseHandler());
183198
return familyForStatusCode(status.getStatusCode()) == Family.SUCCESSFUL;
184199
}
@@ -189,7 +204,10 @@ public void close()
189204
if (!closed.getAndSet(true)) {
190205
URI uri = currentResults.get().getNextUri();
191206
if (uri != null) {
192-
Request request = prepareDelete().setUri(uri).build();
207+
Request request = prepareDelete()
208+
.setHeader(USER_AGENT, USER_AGENT_VALUE)
209+
.setUri(uri)
210+
.build();
193211
httpClient.executeAsync(request, createStatusResponseHandler());
194212
}
195213
}

‎presto-jdbc/src/main/java/com/facebook/presto/jdbc/JdbcConnection.java‎

Lines changed: 19 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.facebook.presto.client.ClientSession;
44
import com.facebook.presto.client.StatementClient;
5+
import com.google.common.base.Objects;
56
import com.google.common.net.HostAndPort;
67

78
import java.net.URI;
@@ -24,13 +25,15 @@
2425
import java.sql.Struct;
2526
import java.util.Map;
2627
import java.util.Properties;
28+
import java.util.concurrent.ConcurrentHashMap;
2729
import java.util.concurrent.Executor;
2830
import java.util.concurrent.atomic.AtomicBoolean;
2931
import java.util.concurrent.atomic.AtomicReference;
3032

3133
import static com.facebook.presto.jdbc.Driver.DRIVER_NAME;
3234
import static com.facebook.presto.jdbc.Driver.DRIVER_VERSION;
3335
import static com.google.common.base.Preconditions.checkNotNull;
36+
import static com.google.common.collect.Maps.fromProperties;
3437
import static io.airlift.http.client.HttpUriBuilder.uriBuilder;
3538

3639
public class JdbcConnection
@@ -42,6 +45,7 @@ public class JdbcConnection
4245
private final URI uri;
4346
private final HostAndPort address;
4447
private final String user;
48+
private final Map<String, String> clientInfo = new ConcurrentHashMap<>();
4549
private final QueryExecutor queryExecutor;
4650

4751
JdbcConnection(URI uri, String user)
@@ -373,28 +377,36 @@ public boolean isValid(int timeout)
373377
public void setClientInfo(String name, String value)
374378
throws SQLClientInfoException
375379
{
376-
throw new UnsupportedOperationException("setClientInfo");
380+
checkNotNull(name, "name is null");
381+
if (value != null) {
382+
clientInfo.put(name, value);
383+
}
384+
else {
385+
clientInfo.remove(value);
386+
}
377387
}
378388

379389
@Override
380390
public void setClientInfo(Properties properties)
381391
throws SQLClientInfoException
382392
{
383-
throw new UnsupportedOperationException("setClientInfo");
393+
clientInfo.putAll(fromProperties(properties));
384394
}
385395

386396
@Override
387397
public String getClientInfo(String name)
388398
throws SQLException
389399
{
390-
throw new UnsupportedOperationException("getClientInfo");
400+
return clientInfo.get(name);
391401
}
392402

393403
@Override
394404
public Properties getClientInfo()
395405
throws SQLException
396406
{
397-
throw new UnsupportedOperationException("getClientInfo");
407+
Properties properties = new Properties();
408+
properties.putAll(clientInfo);
409+
return properties;
398410
}
399411

400412
@Override
@@ -479,7 +491,9 @@ String getUser()
479491
StatementClient startQuery(String sql)
480492
{
481493
URI uri = createHttpUri(address);
482-
ClientSession session = new ClientSession(uri, user, catalog.get(), schema.get(), false);
494+
495+
String source = Objects.firstNonNull(clientInfo.get("ApplicationName"), "presto-jdbc");
496+
ClientSession session = new ClientSession(uri, user, source, catalog.get(), schema.get(), false);
483497
return queryExecutor.startQuery(session, sql);
484498
}
485499

‎presto-main/src/main/java/com/facebook/presto/event/query/QueryCompletionEvent.java‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,8 +20,11 @@ public class QueryCompletionEvent
2020
{
2121
private final QueryId queryId;
2222
private final String user;
23+
private final String source;
2324
private final String catalog;
2425
private final String schema;
26+
private final String remoteClientAddress;
27+
private final String userAgent;
2528
private final QueryState queryState;
2629
private final URI uri;
2730
private final List<String> fieldNames;
@@ -47,9 +50,9 @@ public class QueryCompletionEvent
4750
public QueryCompletionEvent(
4851
QueryId queryId,
4952
String user,
50-
String catalog,
53+
String source, String catalog,
5154
String schema,
52-
QueryState queryState,
55+
String remoteClientAddress, String userAgent, QueryState queryState,
5356
URI uri,
5457
List<String> fieldNames,
5558
String query,
@@ -69,8 +72,11 @@ public QueryCompletionEvent(
6972
{
7073
this.queryId = queryId;
7174
this.user = user;
75+
this.source = source;
7276
this.catalog = catalog;
7377
this.schema = schema;
78+
this.remoteClientAddress = remoteClientAddress;
79+
this.userAgent = userAgent;
7480
this.queryState = queryState;
7581
this.uri = uri;
7682
this.fieldNames = ImmutableList.copyOf(fieldNames);
@@ -120,6 +126,12 @@ public String getUser()
120126
return user;
121127
}
122128

129+
@EventField
130+
public String getSource()
131+
{
132+
return source;
133+
}
134+
123135
@EventField
124136
public String getCatalog()
125137
{
@@ -132,6 +144,18 @@ public String getSchema()
132144
return schema;
133145
}
134146

147+
@EventField
148+
public String getRemoteClientAddress()
149+
{
150+
return remoteClientAddress;
151+
}
152+
153+
@EventField
154+
public String getUserAgent()
155+
{
156+
return userAgent;
157+
}
158+
135159
@EventField
136160
public String getQueryState()
137161
{

‎presto-main/src/main/java/com/facebook/presto/event/query/QueryCreatedEvent.java‎

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,25 +14,31 @@ public class QueryCreatedEvent
1414
{
1515
private final QueryId queryId;
1616
private final String user;
17+
private final String source;
1718
private final String catalog;
1819
private final String schema;
20+
private final String remoteClientAddress;
21+
private final String userAgent;
1922
private final URI uri;
2023
private final String query;
2124
private final DateTime createTime;
2225

2326
public QueryCreatedEvent(
2427
QueryId queryId,
2528
String user,
26-
String catalog,
29+
String source, String catalog,
2730
String schema,
28-
URI uri,
31+
String remoteClientAddress, String userAgent, URI uri,
2932
String query,
3033
DateTime createTime)
3134
{
3235
this.queryId = queryId;
3336
this.user = user;
37+
this.source = source;
3438
this.catalog = catalog;
3539
this.schema = schema;
40+
this.remoteClientAddress = remoteClientAddress;
41+
this.userAgent = userAgent;
3642
this.uri = uri;
3743
this.query = query;
3844
this.createTime = createTime;
@@ -50,6 +56,12 @@ public String getUser()
5056
return user;
5157
}
5258

59+
@EventField
60+
public String getSource()
61+
{
62+
return source;
63+
}
64+
5365
@EventField
5466
public String getCatalog()
5567
{
@@ -62,6 +74,18 @@ public String getSchema()
6274
return schema;
6375
}
6476

77+
@EventField
78+
public String getRemoteClientAddress()
79+
{
80+
return remoteClientAddress;
81+
}
82+
83+
@EventField
84+
public String getUserAgent()
85+
{
86+
return userAgent;
87+
}
88+
6589
@EventField
6690
public String getUri()
6791
{

‎presto-main/src/main/java/com/facebook/presto/event/query/QueryMonitor.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,8 +31,11 @@ public void createdEvent(QueryInfo queryInfo)
3131
new QueryCreatedEvent(
3232
queryInfo.getQueryId(),
3333
queryInfo.getSession().getUser(),
34+
queryInfo.getSession().getSource(),
3435
queryInfo.getSession().getCatalog(),
3536
queryInfo.getSession().getSchema(),
37+
queryInfo.getSession().getRemoteUserAddress(),
38+
queryInfo.getSession().getUserAgent(),
3639
queryInfo.getSelf(),
3740
queryInfo.getQuery(),
3841
queryInfo.getQueryStats().getCreateTime()
@@ -48,8 +51,11 @@ public void completionEvent(QueryInfo queryInfo)
4851
new QueryCompletionEvent(
4952
queryInfo.getQueryId(),
5053
queryInfo.getSession().getUser(),
54+
queryInfo.getSession().getSource(),
5155
queryInfo.getSession().getCatalog(),
5256
queryInfo.getSession().getSchema(),
57+
queryInfo.getSession().getRemoteUserAddress(),
58+
queryInfo.getSession().getUserAgent(),
5359
queryInfo.getState(),
5460
queryInfo.getSelf(),
5561
queryInfo.getFieldNames(),

‎presto-main/src/main/java/com/facebook/presto/importer/PeriodicImportRunnable.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,7 @@ public void doRun()
5959
QualifiedTableName dstTable = job.getDstTable();
6060
String sql = String.format("REFRESH MATERIALIZED VIEW %s", dstTable.getTableName());
6161

62-
ClientSession session = new ClientSession(serverUri(), "periodic-import", dstTable.getCatalogName(), dstTable.getSchemaName(), false);
62+
ClientSession session = new ClientSession(serverUri(), "periodic-import", "periodic-import", dstTable.getCatalogName(), dstTable.getSchemaName(), false);
6363
StatementClient client = new StatementClient(httpClient, queryResultsCodec, session, sql);
6464

6565
// don't delete this line, it is what actually pulls the data from the query...

0 commit comments

Comments
 (0)