/
ConcreteConnectionConfiguration.java
130 lines (111 loc) · 4.94 KB
/
ConcreteConnectionConfiguration.java
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
package io.burt.athena.configuration;
import io.burt.athena.polling.PollingStrategies;
import io.burt.athena.polling.PollingStrategy;
import io.burt.athena.result.PreloadingStandardResult;
import io.burt.athena.result.Result;
import io.burt.athena.result.S3Result;
import io.burt.athena.result.StandardResult;
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.athena.AthenaAsyncClient;
import software.amazon.awssdk.services.athena.model.QueryExecution;
import software.amazon.awssdk.services.s3.S3AsyncClient;
import java.time.Duration;
class ConcreteConnectionConfiguration implements ConnectionConfiguration {
private final Region awsRegion;
private final String databaseName;
private final String workGroupName;
private final String outputLocation;
private final Duration networkTimeout;
private final Duration queryTimeout;
private final ResultLoadingStrategy resultLoadingStrategy;
private AthenaAsyncClient athenaClient;
private S3AsyncClient s3Client;
private PollingStrategy pollingStrategy;
ConcreteConnectionConfiguration(Region awsRegion, String databaseName, String workGroupName, String outputLocation, Duration networkTimeout, Duration queryTimeout, ResultLoadingStrategy resultLoadingStrategy) {
this.awsRegion = awsRegion;
this.databaseName = databaseName;
this.workGroupName = workGroupName;
this.outputLocation = outputLocation;
this.networkTimeout = networkTimeout;
this.queryTimeout = queryTimeout;
this.resultLoadingStrategy = resultLoadingStrategy;
}
private ConcreteConnectionConfiguration(Region awsRegion, String databaseName, String workGroupName, String outputLocation, Duration networkTimeout, Duration queryTimeout, ResultLoadingStrategy resultLoadingStrategy, AthenaAsyncClient athenaClient, S3AsyncClient s3Client, PollingStrategy pollingStrategy) {
this(awsRegion, databaseName, workGroupName, outputLocation, networkTimeout, queryTimeout, resultLoadingStrategy);
this.athenaClient = athenaClient;
this.s3Client = s3Client;
this.pollingStrategy = pollingStrategy;
}
@Override
public String databaseName() {
return databaseName;
}
@Override
public String workGroupName() {
return workGroupName;
}
@Override
public String outputLocation() {
return outputLocation;
}
@Override
public Duration networkTimeout() {
return networkTimeout;
}
@Override
public Duration queryTimeout() { return queryTimeout; }
@Override
public AthenaAsyncClient athenaClient() {
if (athenaClient == null) {
athenaClient = AthenaAsyncClient.builder().region(awsRegion).build();
}
return athenaClient;
}
@Override
public S3AsyncClient s3Client() {
if (s3Client == null) {
s3Client = S3AsyncClient.builder().region(awsRegion).build();
}
return s3Client;
}
@Override
public PollingStrategy pollingStrategy() {
if (pollingStrategy == null) {
pollingStrategy = PollingStrategies.backoff(Duration.ofMillis(10), Duration.ofSeconds(5));
}
return pollingStrategy;
}
@Override
public ConnectionConfiguration withDatabaseName(String databaseName) {
return new ConcreteConnectionConfiguration(awsRegion, databaseName, workGroupName, outputLocation, networkTimeout, queryTimeout, resultLoadingStrategy, athenaClient, s3Client, pollingStrategy);
}
@Override
public ConnectionConfiguration withNetworkTimeout(Duration networkTimeout) {
return new ConcreteConnectionConfiguration(awsRegion, databaseName, workGroupName, outputLocation, networkTimeout, queryTimeout, resultLoadingStrategy, athenaClient, s3Client, pollingStrategy);
}
@Override
public ConnectionConfiguration withQueryTimeout(Duration queryTimeout) {
return new ConcreteConnectionConfiguration(awsRegion, databaseName, workGroupName, outputLocation, networkTimeout, queryTimeout, resultLoadingStrategy, athenaClient, s3Client, pollingStrategy);
}
@Override
public Result createResult(QueryExecution queryExecution) {
if (resultLoadingStrategy == ResultLoadingStrategy.GET_EXECUTION_RESULTS) {
return new PreloadingStandardResult(athenaClient(), queryExecution, StandardResult.MAX_FETCH_SIZE, Duration.ofSeconds(10));
} else if (resultLoadingStrategy == ResultLoadingStrategy.S3) {
return new S3Result(s3Client(), queryExecution, Duration.ofSeconds(10));
} else {
throw new IllegalStateException(String.format("No such result loading strategy: %s", queryExecution));
}
}
@Override
public void close() {
if (athenaClient != null) {
athenaClient.close();
athenaClient = null;
}
if (s3Client != null) {
s3Client.close();
s3Client = null;
}
}
}