AirbyteApiClient2: Enable retries on 5xx responses (#11524)

airbytehq · Mar 6, 2024 · 3888b97 · 3888b97
1 parent 89139e2
commit 3888b97
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/airbyte-api/src/main/kotlin/AirbyteApiClient2.kt b/airbyte-api/src/main/kotlin/AirbyteApiClient2.kt
@@ -24,6 +24,7 @@ import io.airbyte.api.client2.generated.StateApi
 import io.airbyte.api.client2.generated.StreamStatusesApi
 import io.airbyte.api.client2.generated.WorkspaceApi
 import okhttp3.OkHttpClient
+import java.io.IOException
 
 /**
  * This class wraps all the generated API clients and provides a single entry point. This class is meant
@@ -55,8 +56,15 @@ class AirbyteApiClient2
   constructor(
     basePath: String,
     policy: RetryPolicy<okhttp3.Response> = RetryPolicy.ofDefaults(),
-    httpClient: OkHttpClient = OkHttpClient(),
+    var httpClient: OkHttpClient = OkHttpClient(),
+    throwOn5xx: Boolean = true,
   ) {
+    init {
+      if (throwOn5xx) {
+        httpClient = httpClient.newBuilder().addInterceptor(ThrowOn5xxInterceptor()).build()
+      }
+    }
+
     val connectionApi = ConnectionApi(basePath = basePath, client = httpClient, policy = policy)
     val connectorBuilderProjectApi = ConnectorBuilderProjectApi(basePath = basePath, client = httpClient, policy = policy)
     val deploymentMetadataApi = DeploymentMetadataApi(basePath = basePath, client = httpClient, policy = policy)
@@ -76,3 +84,13 @@ class AirbyteApiClient2
     val streamStatusesApi = StreamStatusesApi(basePath = basePath, client = httpClient, policy = policy)
     val secretPersistenceConfigApi = SecretsPersistenceConfigApi(basePath = basePath, client = httpClient, policy = policy)
   }
+
+class ThrowOn5xxInterceptor : okhttp3.Interceptor {
+  override fun intercept(chain: okhttp3.Interceptor.Chain): okhttp3.Response {
+    val response = chain.proceed(chain.request())
+    if (response.code >= 500) {
+      throw IOException("HTTP error: ${response.code} ${response.message}")
+    }
+    return response
+  }
+}
diff --git a/...-tests/src/test-acceptance/java/io/airbyte/test/acceptance/VersioningAcceptanceTests.java b/...-tests/src/test-acceptance/java/io/airbyte/test/acceptance/VersioningAcceptanceTests.java
@@ -44,7 +44,7 @@ static void init() throws IOException, URISyntaxException {
         .withBackoff(Duration.ofSeconds(1), Duration.ofSeconds(60)).build();
 
     final OkHttpClient client = new OkHttpClient.Builder().readTimeout(Duration.ofSeconds(60)).build();
-    apiClient2 = new AirbyteApiClient2(String.format("%s/api", AIRBYTE_SERVER_HOST), policy, client);
+    apiClient2 = new AirbyteApiClient2(String.format("%s/api", AIRBYTE_SERVER_HOST), policy, client, /* throwOn5xx */ true);
 
     workspaceId = apiClient2.getWorkspaceApi().listWorkspaces().getWorkspaces().get(0).getWorkspaceId();
   }