feat: configurable http codes to new access token (#290)

Author: 1DIce

src/main/java/io/aiven/kafka/connect/http/config/HttpSinkConfig.java

@@ -65,6 +65,7 @@ public final class HttpSinkConfig extends AbstractConfig {
     private static final String OAUTH2_CLIENT_AUTHORIZATION_MODE_CONFIG = "oauth2.client.authorization.mode";
     private static final String OAUTH2_CLIENT_SCOPE_CONFIG = "oauth2.client.scope";
     private static final String OAUTH2_RESPONSE_TOKEN_PROPERTY_CONFIG = "oauth2.response.token.property";
+    private static final String OAUTH2_TOKEN_RENEW_ON_STATUS_CODES_CONFIG ="oauth2.token.renew.on.status.codes";
 
     private static final String BATCHING_GROUP = "Batching";
     private static final String BATCHING_ENABLED_CONFIG = "batching.enabled";
@@ -448,6 +449,48 @@ public String toString() {
                 List.of(OAUTH2_ACCESS_TOKEN_URL_CONFIG, OAUTH2_CLIENT_ID_CONFIG, OAUTH2_CLIENT_SECRET_CONFIG,
                         OAUTH2_CLIENT_AUTHORIZATION_MODE_CONFIG, OAUTH2_CLIENT_SCOPE_CONFIG)
         );
+        configDef.define(
+                OAUTH2_TOKEN_RENEW_ON_STATUS_CODES_CONFIG,
+                ConfigDef.Type.LIST,
+                "401",
+                new ConfigDef.Validator() {
+
+                    @Override
+                    public void ensureValid(String name, Object value) {
+                        if (value == null) {
+                            throw new ConfigException(name, null, "can't be null");
+                        }
+                        if (!(value instanceof List)) {
+                            throw new ConfigException(name, value, "must be a list");
+                        }
+                        for (Object entry : (List) value) {
+                            if (!(entry instanceof String) || !isNumber((String) entry)) {
+                                throw new ConfigException(name, value, "must be a list of numbers");
+                            }
+                        }
+                    }
+
+                    private boolean isNumber(String value) {
+                        try {
+                            Integer.parseInt(value);
+                            return true;
+                        } catch (NumberFormatException e) {
+                            return false;
+                        }
+                    }
+
+                    @Override
+                    public String toString() {
+                        return "List of HTTP status codes";
+                    }
+                },
+                ConfigDef.Importance.LOW,
+                "Comma separated list of HTTP response status codes which should trigger an access token renewal",
+                CONNECTION_GROUP,
+                groupCounter++,
+                ConfigDef.Width.LONG,
+                OAUTH2_TOKEN_RENEW_ON_STATUS_CODES_CONFIG
+        );
     }
 
     private static void addBatchingConfigGroup(final ConfigDef configDef) {
@@ -819,6 +862,10 @@ public final String oauth2ResponseTokenProperty() {
         return getString(OAUTH2_RESPONSE_TOKEN_PROPERTY_CONFIG);
     }
 
+    public final List<Integer> oauth2RenewTokenOnStatusCodes() {
+        return getList(OAUTH2_TOKEN_RENEW_ON_STATUS_CODES_CONFIG).stream().map(Integer::valueOf).collect(Collectors.toList());
+    }
+
     public final boolean hasProxy() {
         return getString(HTTP_PROXY_HOST) != null;
     }

src/main/java/io/aiven/kafka/connect/http/sender/OAuth2HttpSender.java

@@ -48,10 +48,10 @@ class OAuth2HttpSender extends AbstractHttpSender implements HttpSender {
     protected HttpResponse<String> sendWithRetries(
         final Builder requestBuilder, final HttpResponseHandler originHttpResponseHandler, final int retries
     ) {
-        // This handler allows to request a new access token if a 401 occurs, meaning the session might be expired
+        // This handler allows to request a new access token if a an expected error occurs, meaning the session might be expired
         final HttpResponseHandler handler = (response, remainingRetries) -> {
-            // If the response has a 401 error and we have retries left, we attempt to renew the session
-            if (response.statusCode() == 401 && remainingRetries > 0) {
+            // If the response has one of the configured error codes and we have retries left, we attempt to renew the session
+            if (config.oauth2RenewTokenOnStatusCodes().contains(response.statusCode()) && remainingRetries > 0) {
                 // Update the request builder with the new access token
                 ((OAuth2AuthHttpRequestBuilder) this.httpRequestBuilder).renewAccessToken(requestBuilder);
                 // Retry the call and decrease the retries counter to avoid looping on token renewal

src/test/java/io/aiven/kafka/connect/http/config/HttpSinkConfigTest.java

@@ -21,6 +21,7 @@
 import java.net.URISyntaxException;
 import java.util.Collections;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import java.util.stream.Stream;
@@ -77,7 +78,9 @@ void correctMinimalConfig() throws URISyntaxException {
                 .returns("access_token", from(HttpSinkConfig::oauth2ResponseTokenProperty))
                 .returns(null, from(HttpSinkConfig::kafkaRetryBackoffMs))
                 .returns(false, from(HttpSinkConfig::hasProxy))
-                .returns(false, from(HttpSinkConfig::sslTrustAllCertificates));
+                .returns(false, from(HttpSinkConfig::sslTrustAllCertificates))
+                .returns(List.of(401), from(HttpSinkConfig::oauth2RenewTokenOnStatusCodes));
+
     }
 
     @Test
@@ -195,7 +198,8 @@ void validOAuth2FullConfiguration() throws URISyntaxException {
                 "oauth2.client.secret", "client_secret",
                 "oauth2.client.authorization.mode", "url",
                 "oauth2.client.scope", "scope1,scope2",
-                "oauth2.response.token.property", "moooooo"
+                "oauth2.response.token.property", "moooooo",
+                "oauth2.token.renew.on.status.codes", "401, 403"
         );
 
         final var config = new HttpSinkConfig(oauth2Config);
@@ -206,7 +210,22 @@ void validOAuth2FullConfiguration() throws URISyntaxException {
                 .returns("client_secret", from(httpSinkConfig -> httpSinkConfig.oauth2ClientSecret().value()))
                 .returns("scope1,scope2", from(HttpSinkConfig::oauth2ClientScope))
                 .returns(OAuth2AuthorizationMode.URL, from(HttpSinkConfig::oauth2AuthorizationMode))
-                .returns("moooooo", from(HttpSinkConfig::oauth2ResponseTokenProperty));
+                .returns("moooooo", from(HttpSinkConfig::oauth2ResponseTokenProperty))
+                .returns(List.of(401, 403), from(HttpSinkConfig::oauth2RenewTokenOnStatusCodes));
+    }
+
+    @Test
+    void invalidHttpErrorCode() {
+        final Map<String, String> properties = Map.of(
+                "http.url", "http://localhost:8090",
+                "http.authorization.type", "none",
+                "oauth2.token.renew.on.status.codes", "wrong"
+        );
+
+        assertThatExceptionOfType(ConfigException.class)
+                .describedAs("Expected config exception due to malformed status code list")
+                .isThrownBy(() -> new HttpSinkConfig(properties))
+                .withMessage("Invalid value [wrong] for configuration oauth2.token.renew.on.status.codes: must be a list of numbers");
     }
 
     @Test

src/test/java/io/aiven/kafka/connect/http/sender/OAuth2HttpSenderTest.java

@@ -34,6 +34,9 @@
 import org.assertj.core.api.InstanceOfAssertFactories;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.CsvSource;
+import org.junit.jupiter.params.provider.ValueSource;
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.Mockito;
@@ -228,8 +231,13 @@ void reuseAccessToken() throws Exception {
 
     }
 
-    @Test
-    void refreshAccessToken() throws Exception {
+    @ParameterizedTest
+    @CsvSource( {
+            "'401'    , '401'",
+            "'401,403', '401'",
+            "'401,403', '403'"
+    })
+    void refreshAccessToken(final String configuredErrorCodes, final String sendErrorCode) throws Exception {
 
         // first call to retrieve an access token
         final HttpResponse<String> mockedAccessTokenResponse = mock(HttpResponse.class);
@@ -241,15 +249,17 @@ void refreshAccessToken() throws Exception {
         when(oauth2AccessTokenHttpSender.call()).thenReturn(
             mockedAccessTokenResponse, mockedAccessTokenResponseRefreshed);
 
-        // Mock a 2nd response with 401.
+        // Mock a 2nd response with error code.
         final HttpResponse<String> errorResponse = mock(HttpResponse.class);
-        when(errorResponse.statusCode()).thenReturn(401);
-        // Mock a 2nd response with 401.
+        when(errorResponse.statusCode()).thenReturn(Integer.parseInt(sendErrorCode));
+        // Mock a 2nd response with 200.
         final HttpResponse<String> normalResponse = mock(HttpResponse.class);
         when(normalResponse.statusCode()).thenReturn(200);
 
         // Build the configuration
-        final HttpSinkConfig config = new HttpSinkConfig(defaultConfig());
+        Map<String,String> configMap =new HashMap<>(defaultConfig());
+        configMap.put("oauth2.token.renew.on.status.codes", configuredErrorCodes);
+        final HttpSinkConfig config = new HttpSinkConfig(configMap);
 
         // Mock the Client and Response
         when(mockedClient.send(any(HttpRequest.class), any(BodyHandler.class))).thenReturn(mockedResponse,