From a2b8463247567347ebd13ab5f0763b383e76f3bf Mon Sep 17 00:00:00 2001
From: Igor Motov <igor@motovs.org>
Date: Mon, 30 Jul 2018 16:31:02 -0400
Subject: [PATCH] SQL: Add test for handling of partial results

Makes sure that partial results are rejected in the SQL mode.

Closes #32284
---
 .../qa/sql/nosecurity/JdbcShardFailureIT.java | 11 ++++
 .../qa/sql/jdbc/JdbcShardFailureTestCase.java | 54 +++++++++++++++++++
 2 files changed, 65 insertions(+)
 create mode 100644 x-pack/qa/sql/no-security/src/test/java/org/elasticsearch/xpack/qa/sql/nosecurity/JdbcShardFailureIT.java
 create mode 100644 x-pack/qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/jdbc/JdbcShardFailureTestCase.java

diff --git a/x-pack/qa/sql/no-security/src/test/java/org/elasticsearch/xpack/qa/sql/nosecurity/JdbcShardFailureIT.java b/x-pack/qa/sql/no-security/src/test/java/org/elasticsearch/xpack/qa/sql/nosecurity/JdbcShardFailureIT.java
new file mode 100644
index 0000000000000..4580533c4c4c0
--- /dev/null
+++ b/x-pack/qa/sql/no-security/src/test/java/org/elasticsearch/xpack/qa/sql/nosecurity/JdbcShardFailureIT.java
@@ -0,0 +1,11 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.xpack.qa.sql.nosecurity;
+
+import org.elasticsearch.xpack.qa.sql.jdbc.JdbcShardFailureTestCase;
+
+public class JdbcShardFailureIT extends JdbcShardFailureTestCase {
+}
diff --git a/x-pack/qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/jdbc/JdbcShardFailureTestCase.java b/x-pack/qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/jdbc/JdbcShardFailureTestCase.java
new file mode 100644
index 0000000000000..445559372a15e
--- /dev/null
+++ b/x-pack/qa/sql/src/main/java/org/elasticsearch/xpack/qa/sql/jdbc/JdbcShardFailureTestCase.java
@@ -0,0 +1,54 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License;
+ * you may not use this file except in compliance with the Elastic License.
+ */
+package org.elasticsearch.xpack.qa.sql.jdbc;
+
+import org.elasticsearch.client.Request;
+import org.junit.Before;
+
+import java.io.IOException;
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.Statement;
+
+import static org.hamcrest.Matchers.containsString;
+
+/**
+ * Tests for handling of partial results in SQL layer
+ */
+public class JdbcShardFailureTestCase extends JdbcIntegrationTestCase {
+    @Before
+    public void createTestIndex() throws IOException {
+        Request createTest1 = new Request("PUT", "/test1");
+        String body1 = "{\"aliases\":{\"test\":{}}, \"mappings\": {\"doc\": {\"properties\": {\"test_field\":{\"type\":\"integer\"}}}}}";
+        createTest1.setJsonEntity(body1);
+        client().performRequest(createTest1);
+
+        Request createTest2 = new Request("PUT", "/test2");
+        String body2 = "{\"aliases\":{\"test\":{}}, \"mappings\": {\"doc\": {\"properties\": {\"test_field\":{\"type\":\"integer\"}}}}," +
+            "\"settings\": {\"index.routing.allocation.include.node\": \"nowhere\"}}";
+        createTest2.setJsonEntity(body2);
+        createTest2.addParameter("timeout", "100ms");
+        client().performRequest(createTest2);
+
+        Request request = new Request("PUT", "/test1/doc/_bulk");
+        request.addParameter("refresh", "true");
+        StringBuilder bulk = new StringBuilder();
+        for (int i = 0; i < 20; i++) {
+            bulk.append("{\"index\":{}}\n");
+            bulk.append("{\"test_field\":" + i + "}\n");
+        }
+        request.setJsonEntity(bulk.toString());
+        client().performRequest(request);
+    }
+
+
+    public void testPartialResponseHandling() throws SQLException {
+        try (Connection c = esJdbc(); Statement s = c.createStatement()) {
+            SQLException exception = expectThrows(SQLException.class, () -> s.executeQuery("SELECT * FROM test ORDER BY test_field ASC"));
+            assertThat(exception.getMessage(), containsString("Search rejected due to missing shards"));
+        }
+    }
+}