|
21 | 21 | import java.time.ZonedDateTime; |
22 | 22 | import java.time.format.DateTimeFormatter; |
23 | 23 | import java.util.*; |
| 24 | +import java.util.function.Function; |
24 | 25 | import org.slf4j.Logger; |
25 | 26 | import org.slf4j.LoggerFactory; |
26 | 27 |
|
|
29 | 30 | * guessing |
30 | 31 | */ |
31 | 32 | public class ApiClient { |
| 33 | + public static class Builder { |
| 34 | + private Timer timer; |
| 35 | + private Function<Void, Map<String, String>> authenticateFunc; |
| 36 | + private Function<Void, String> getHostFunc; |
| 37 | + private Function<Void, String> getAuthTypeFunc; |
| 38 | + private Integer debugTruncateBytes; |
| 39 | + private HttpClient httpClient; |
| 40 | + private String accountId; |
| 41 | + private RetryStrategyPicker retryStrategyPicker; |
| 42 | + private boolean isDebugHeaders; |
| 43 | + |
| 44 | + public Builder withDatabricksConfig(DatabricksConfig config) { |
| 45 | + this.authenticateFunc = v -> config.authenticate(); |
| 46 | + this.getHostFunc = v -> config.getHost(); |
| 47 | + this.getAuthTypeFunc = v -> config.getAuthType(); |
| 48 | + this.httpClient = config.getHttpClient(); |
| 49 | + this.debugTruncateBytes = config.getDebugTruncateBytes(); |
| 50 | + this.accountId = config.getAccountId(); |
| 51 | + this.retryStrategyPicker = new RequestBasedRetryStrategyPicker(config.getHost()); |
| 52 | + this.isDebugHeaders = config.isDebugHeaders(); |
| 53 | + |
| 54 | + return this; |
| 55 | + } |
| 56 | + |
| 57 | + public Builder withTimer(Timer timer) { |
| 58 | + this.timer = timer; |
| 59 | + return this; |
| 60 | + } |
| 61 | + |
| 62 | + public Builder withAuthenticateFunc(Function<Void, Map<String, String>> authenticateFunc) { |
| 63 | + this.authenticateFunc = authenticateFunc; |
| 64 | + return this; |
| 65 | + } |
| 66 | + |
| 67 | + public Builder withGetHostFunc(Function<Void, String> getHostFunc) { |
| 68 | + this.getHostFunc = getHostFunc; |
| 69 | + return this; |
| 70 | + } |
| 71 | + |
| 72 | + public Builder withGetAuthTypeFunc(Function<Void, String> getAuthTypeFunc) { |
| 73 | + this.getAuthTypeFunc = getAuthTypeFunc; |
| 74 | + return this; |
| 75 | + } |
| 76 | + |
| 77 | + public Builder withHttpClient(HttpClient httpClient) { |
| 78 | + this.httpClient = httpClient; |
| 79 | + return this; |
| 80 | + } |
| 81 | + |
| 82 | + public Builder withRetryStrategyPicker(RetryStrategyPicker retryStrategyPicker) { |
| 83 | + this.retryStrategyPicker = retryStrategyPicker; |
| 84 | + return this; |
| 85 | + } |
| 86 | + |
| 87 | + public ApiClient build() { |
| 88 | + return new ApiClient(this); |
| 89 | + } |
| 90 | + } |
| 91 | + |
32 | 92 | private static final Logger LOG = LoggerFactory.getLogger(ApiClient.class); |
33 | 93 |
|
34 | 94 | private final int maxAttempts; |
35 | 95 |
|
36 | 96 | private final ObjectMapper mapper; |
37 | 97 |
|
38 | | - private final DatabricksConfig config; |
39 | | - |
40 | 98 | private final Random random; |
41 | 99 |
|
42 | 100 | private final HttpClient httpClient; |
43 | 101 | private final BodyLogger bodyLogger; |
44 | 102 | private final RetryStrategyPicker retryStrategyPicker; |
45 | 103 | private final Timer timer; |
| 104 | + private final Function<Void, Map<String, String>> authenticateFunc; |
| 105 | + private final Function<Void, String> getHostFunc; |
| 106 | + private final Function<Void, String> getAuthTypeFunc; |
| 107 | + private final String accountId; |
| 108 | + private final boolean isDebugHeaders; |
46 | 109 | private static final String RETRY_AFTER_HEADER = "retry-after"; |
47 | 110 |
|
48 | 111 | public ApiClient() { |
49 | 112 | this(ConfigLoader.getDefault()); |
50 | 113 | } |
51 | 114 |
|
52 | 115 | public String configuredAccountID() { |
53 | | - return config.getAccountId(); |
| 116 | + return accountId; |
54 | 117 | } |
55 | 118 |
|
56 | 119 | public ApiClient(DatabricksConfig config) { |
57 | 120 | this(config, new SystemTimer()); |
58 | 121 | } |
59 | 122 |
|
60 | 123 | public ApiClient(DatabricksConfig config, Timer timer) { |
61 | | - this.config = config; |
62 | | - config.resolve(); |
63 | | - |
64 | | - Integer rateLimit = config.getRateLimit(); |
65 | | - if (rateLimit == null) { |
66 | | - rateLimit = 15; |
67 | | - } |
68 | | - |
69 | | - Integer debugTruncateBytes = config.getDebugTruncateBytes(); |
| 124 | + this(new Builder().withDatabricksConfig(config.resolve()).withTimer(timer)); |
| 125 | + } |
| 126 | + |
| 127 | + private ApiClient(Builder builder) { |
| 128 | + this.timer = builder.timer != null ? builder.timer : new SystemTimer(); |
| 129 | + this.authenticateFunc = |
| 130 | + builder.authenticateFunc != null |
| 131 | + ? builder.authenticateFunc |
| 132 | + : v -> new HashMap<String, String>(); |
| 133 | + this.getHostFunc = builder.getHostFunc != null ? builder.getHostFunc : v -> ""; |
| 134 | + this.getAuthTypeFunc = builder.getAuthTypeFunc != null ? builder.getAuthTypeFunc : v -> ""; |
| 135 | + this.httpClient = builder.httpClient; |
| 136 | + this.accountId = builder.accountId; |
| 137 | + this.retryStrategyPicker = |
| 138 | + builder.retryStrategyPicker != null |
| 139 | + ? builder.retryStrategyPicker |
| 140 | + : new RequestBasedRetryStrategyPicker(this.getHostFunc.apply(null)); |
| 141 | + this.isDebugHeaders = builder.isDebugHeaders; |
| 142 | + |
| 143 | + Integer debugTruncateBytes = builder.debugTruncateBytes; |
70 | 144 | if (debugTruncateBytes == null) { |
71 | 145 | debugTruncateBytes = 96; |
72 | 146 | } |
73 | 147 |
|
74 | 148 | maxAttempts = 4; |
75 | 149 | mapper = SerDeUtils.createMapper(); |
76 | 150 | random = new Random(); |
77 | | - httpClient = config.getHttpClient(); |
78 | 151 | bodyLogger = new BodyLogger(mapper, 1024, debugTruncateBytes); |
79 | | - retryStrategyPicker = new RequestBasedRetryStrategyPicker(this.config); |
80 | | - this.timer = timer; |
81 | 152 | } |
82 | 153 |
|
83 | 154 | private static <I> void setQuery(Request in, I entity) { |
@@ -203,7 +274,7 @@ private <I> Request prepareBaseRequest(String method, String path, I in) |
203 | 274 | InputStream body = (InputStream) in; |
204 | 275 | return new Request(method, path, body); |
205 | 276 | } else { |
206 | | - String body = serialize(in); |
| 277 | + String body = (in instanceof String) ? (String) in : serialize(in); |
207 | 278 | return new Request(method, path, body); |
208 | 279 | } |
209 | 280 | } |
@@ -245,15 +316,19 @@ private Response executeInner(Request in, String path) { |
245 | 316 | Response out = null; |
246 | 317 |
|
247 | 318 | // Authenticate the request. Failures should not be retried. |
248 | | - in.withHeaders(config.authenticate()); |
| 319 | + in.withHeaders(authenticateFunc.apply(null)); |
249 | 320 |
|
250 | 321 | // Prepend host to URL only after config.authenticate(). |
251 | 322 | // This call may configure the host (e.g. in case of notebook native auth). |
252 | | - in.withUrl(config.getHost() + path); |
| 323 | + in.withUrl(getHostFunc.apply(null) + path); |
253 | 324 |
|
254 | 325 | // Set User-Agent with auth type info, which is available only |
255 | 326 | // after the first invocation to config.authenticate() |
256 | | - String userAgent = String.format("%s auth/%s", UserAgent.asString(), config.getAuthType()); |
| 327 | + String userAgent = UserAgent.asString(); |
| 328 | + String authType = getAuthTypeFunc.apply(null); |
| 329 | + if (authType != "") { |
| 330 | + userAgent += String.format(" auth/%s", authType); |
| 331 | + } |
257 | 332 | in.withHeader("User-Agent", userAgent); |
258 | 333 |
|
259 | 334 | // Make the request, catching any exceptions, as we may want to retry. |
@@ -347,9 +422,9 @@ private String makeLogRecord(Request in, Response out) { |
347 | 422 | StringBuilder sb = new StringBuilder(); |
348 | 423 | sb.append("> "); |
349 | 424 | sb.append(in.getRequestLine()); |
350 | | - if (config.isDebugHeaders()) { |
| 425 | + if (this.isDebugHeaders) { |
351 | 426 | sb.append("\n * Host: "); |
352 | | - sb.append(config.getHost()); |
| 427 | + sb.append(this.getHostFunc.apply(null)); |
353 | 428 | in.getHeaders() |
354 | 429 | .forEach((header, value) -> sb.append(String.format("\n * %s: %s", header, value))); |
355 | 430 | } |
|
0 commit comments