Rest request connection closed retries

src/test/java/net/snowflake/client/jdbc/RestRequestTest.java

@@ -16,6 +16,7 @@
 import static org.mockito.Mockito.when;
 
 import java.io.IOException;
+import java.net.SocketException;
 import java.net.SocketTimeoutException;
 import java.util.ArrayList;
 import java.util.List;
@@ -492,6 +493,29 @@ public CloseableHttpResponse answer(InvocationOnMock invocation) throws Throwabl
     }
   }
 
+  @Test
+  public void testConnectionClosedRetriesSuccessful() throws IOException, SnowflakeSQLException {
+    CloseableHttpClient client = mock(CloseableHttpClient.class);
+    when(client.execute(any(HttpUriRequest.class)))
+            .then(
+                    new Answer<CloseableHttpResponse>() {
+                      int callCount = 0;
+                      @Override
+                      public CloseableHttpResponse answer(InvocationOnMock invocationOnMock) throws Throwable {
+                        callCount += 1;
+                        if (callCount >= 1) {
+                          return successResponse();
+                        }
+                        else {
+                          throw new SocketException("Connection reset");
+                        }
+                      }
+                    }
+            );
+
+    execute(client, "fakeurl.com/?requestId=abcd-1234", 0, 0, 0, true, false, 1);
+  }
+
   @Test(expected = SnowflakeSQLException.class)
   public void testLoginMaxRetries() throws IOException, SnowflakeSQLException {
     boolean telemetryEnabled = TelemetryService.getInstance().isEnabled();

src/test/java/net/snowflake/client/jdbc/RestRequestWiremockTest.java

@@ -0,0 +1,250 @@
+package net.snowflake.client.jdbc;
+
+import net.snowflake.client.RunningNotOnGithubActionsMac;
+import net.snowflake.client.RunningNotOnJava21;
+import net.snowflake.client.RunningNotOnJava8;
+import net.snowflake.client.core.ExecTimeTelemetryData;
+import net.snowflake.client.log.SFLogger;
+import net.snowflake.client.log.SFLoggerFactory;
+import org.apache.http.ConnectionClosedException;
+import org.apache.http.client.methods.CloseableHttpResponse;
+import org.apache.http.client.methods.HttpGet;
+import org.apache.http.client.methods.HttpPost;
+import org.apache.http.entity.StringEntity;
+import org.apache.http.impl.client.CloseableHttpClient;
+import org.apache.http.impl.client.HttpClients;
+import org.apache.http.util.EntityUtils;
+import org.junit.AfterClass;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+import java.io.File;
+import java.io.IOException;
+import java.net.ServerSocket;
+import java.nio.file.Paths;
+import java.sql.SQLException;
+import java.time.Duration;
+import java.util.Objects;
+import java.util.Properties;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import static net.snowflake.client.jdbc.SnowflakeUtil.systemGetProperty;
+import static org.awaitility.Awaitility.await;
+import static org.junit.Assume.*;
+
+public class RestRequestWiremockTest {
+    private static final SFLogger logger = SFLoggerFactory.getLogger(ProxyLatestIT.class);
+    private static final String WIREMOCK_HOME_DIR = ".wiremock";
+    private static final String WIREMOCK_M2_PATH =
+            "/.m2/repository/org/wiremock/wiremock-standalone/3.8.0/wiremock-standalone-3.8.0.jar";
+    private static final String WIREMOCK_HOST = "localhost";
+    private static final String TRUST_STORE_PROPERTY = "javax.net.ssl.trustStore";
+    private static int httpProxyPort;
+    private static int httpsProxyPort;
+    private static String originalTrustStorePath;
+    private static Process wiremockStandalone;
+
+    private static void startWiremockStandAlone() {
+        // retrying in case of fail in port bindings
+        await()
+                .alias("wait for wiremock responding")
+                .atMost(Duration.ofSeconds(10))
+                .until(
+                        () -> {
+                            try {
+                                httpProxyPort = findFreePort();
+                                httpsProxyPort = findFreePort();
+                                wiremockStandalone =
+                                        new ProcessBuilder(
+                                                "java",
+                                                "-jar",
+                                                getWiremockStandAlonePath(),
+                                                "--root-dir",
+                                                System.getProperty("user.dir")
+                                                        + File.separator
+                                                        + WIREMOCK_HOME_DIR
+                                                        + File.separator,
+                                                "--enable-browser-proxying", // work as forward proxy
+                                                "--proxy-pass-through",
+                                                "false", // pass through only matched requests
+                                                "--port",
+                                                String.valueOf(httpProxyPort),
+                                                "--https-port",
+                                                String.valueOf(httpsProxyPort),
+                                                "--https-keystore",
+                                                getResourceURL("wiremock" + File.separator + "ca-cert.jks"),
+                                                "--ca-keystore",
+                                                getResourceURL("wiremock" + File.separator + "ca-cert.jks"))
+                                                .inheritIO()
+                                                .start();
+                                waitForWiremock();
+                                return true;
+                            } catch (Exception e) {
+                                logger.warn("Failed to start wiremock, retrying: ", e);
+                                return false;
+                            }
+                        });
+    }
+
+    private static int findFreePort() {
+        try {
+            ServerSocket socket = new ServerSocket(0);
+            int port = socket.getLocalPort();
+            socket.close();
+            return port;
+        } catch (Exception e) {
+            throw new RuntimeException(e);
+        }
+    }
+
+    private static String getWiremockStandAlonePath() {
+        return System.getProperty("user.home") + WIREMOCK_M2_PATH;
+    }
+
+    private static String getResourceURL(String relativePath) {
+        return Paths.get(systemGetProperty("user.dir"), "src", "test", "resources", relativePath)
+                .toAbsolutePath()
+                .toString();
+    }
+
+    private static void waitForWiremock() {
+        await()
+                .pollDelay(Duration.ofSeconds(1))
+                .atMost(Duration.ofSeconds(3))
+                .until(RestRequestWiremockTest::isWiremockResponding);
+    }
+
+    private static boolean isWiremockResponding() {
+        try (CloseableHttpClient httpClient = HttpClients.createDefault()) {
+            HttpGet request =
+                    new HttpGet(String.format("http://%s:%d/__admin/mappings", WIREMOCK_HOST, httpProxyPort));
+            CloseableHttpResponse response = httpClient.execute(request);
+            return response.getStatusLine().getStatusCode() == 200;
+        } catch (Exception e) {
+            logger.warn("Waiting for wiremock to respond: ", e);
+        }
+        return false;
+    }
+
+    private static void stopWiremockStandAlone() {
+        if (wiremockStandalone != null) {
+            wiremockStandalone.destroyForcibly();
+            await()
+                    .alias("stop wiremock")
+                    .atMost(Duration.ofSeconds(10))
+                    .until(() -> !wiremockStandalone.isAlive());
+        }
+    }
+
+    @BeforeClass
+    public static void setUpClass() {
+        assumeFalse(RunningNotOnJava8.isRunningOnJava8());
+        assumeFalse(RunningNotOnJava21.isRunningOnJava21());
+        assumeFalse(
+                RunningNotOnGithubActionsMac
+                        .isRunningOnGithubActionsMac()); // disabled until issue with access to localhost
+        // (https://github.com/snowflakedb/snowflake-jdbc/pull/1807#discussion_r1686229430) is fixed on
+        // github actions mac image. Ticket to enable when fixed: SNOW-1555950
+        originalTrustStorePath = systemGetProperty(TRUST_STORE_PROPERTY);
+        startWiremockStandAlone();
+    }
+
+    private static void addMapping(String mapping) {
+        try (CloseableHttpClient httpClient = HttpClients.createDefault()) {
+            HttpPost request =
+                new HttpPost(String.format("http://%s:%d/__admin/mappings/import", WIREMOCK_HOST, httpProxyPort));
+
+            request.setEntity(new StringEntity(mapping));
+            CloseableHttpResponse response = httpClient.execute(request);
+            if (response.getStatusLine().getStatusCode() != 200) {
+                System.out.println(response.getStatusLine().getStatusCode());
+                System.out.println(EntityUtils.toString(response.getEntity()));
+            }
+            assumeTrue( response.getStatusLine().getStatusCode() == 200);
+        } catch (Exception e) {
+            e.printStackTrace();
+            assumeNoException(e);
+        }
+    }
+
+    String example_mapping = "" +
+            "{\n" +
+            "  \"request\": {\n" +
+            "    \"method\": \"GET\",\n" +
+            "    \"url\": \"/some/thing\"\n" +
+            "  },\n" +
+            "\n" +
+            "  \"response\": {\n" +
+            "    \"status\": 200,\n" +
+            "    \"body\": \"Hello, world!\",\n" +
+            "    \"headers\": {\n" +
+            "        \"Content-Type\": \"text/plain\"\n" +
+            "    }\n" +
+            "  }\n" +
+            "}";
+
+    String connectionResetByPeerScenario = "{\n" +
+            "    \"mappings\": [\n" +
+            "        {\n" +
+            "            \"scenarioName\": \"Connection reset by peer\",\n" +
+            "            \"requiredScenarioState\": \"Started\",\n" +
+            "            \"newScenarioState\": \"Connection is stable\",\n" +
+            "            \"request\": {\n" +
+            "                \"method\": \"GET\",\n" +
+            "                \"url\": \"/endpoint\"\n" +
+            "            },\n" +
+            "            \"response\": {\n" +
+            "                \"fault\": \"CONNECTION_RESET_BY_PEER\"\n" +
+            "            }\n" +
+            "        },\n" +
+            "        {\n" +
+            "            \"scenarioName\": \"Connection reset by peer\",\n" +
+            "            \"requiredScenarioState\": \"Connection is stable\",\n" +
+            "            \"request\": {\n" +
+            "                \"method\": \"GET\",\n" +
+            "                \"url\": \"/endpoint\"\n" +
+            "            },\n" +
+            "            \"response\": {\n" +
+            "                \"status\": 200\n" +
+            "            }\n" +
+            "        }\n" +
+            "    ],\n" +
+            "    \"importOptions\": {\n" +
+            "        \"duplicatePolicy\": \"IGNORE\",\n" +
+            "        \"deleteAllNotInImport\": true\n" +
+            "    }" +
+            "}";
+
+    @Test
+    public void testProxyIsUsedWhenSetInProperties() throws Exception {
+        addMapping(connectionResetByPeerScenario);
+        try (CloseableHttpClient httpClient = HttpClients.createDefault()) {
+            HttpGet request =
+                new HttpGet(String.format("http://%s:%d/endpoint", WIREMOCK_HOST, httpProxyPort));
+            RestRequest.execute(
+                    httpClient,
+                    request,
+                    0,
+                    0,
+                    0,
+                    1,
+                    0,
+                    new AtomicBoolean(false),
+                    false,
+                    true,
+                    false,
+                    true,
+                    new ExecTimeTelemetryData());
+
+            CloseableHttpResponse response = httpClient.execute(request);
+            assert (response.getStatusLine().getStatusCode() == 200);
+        } catch (Exception e) {
+            throw e;
+        }
+    }
+
+    @AfterClass
+    public static void tearDownClass() {
+        stopWiremockStandAlone();
+    }
+}