Consider ThriftServerQueryTestSuite

apache · Mar 28, 2020 · 8eac8d6 · 8eac8d6
1 parent 0511d99
commit 8eac8d6
Showing 1 changed file with 66 additions and 53 deletions.
diff --git a/...er/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerQueryTestSuite.scala b/...er/src/test/scala/org/apache/spark/sql/hive/thriftserver/ThriftServerQueryTestSuite.scala
@@ -76,11 +76,12 @@ class ThriftServerQueryTestSuite extends SQLQueryTestSuite with SharedThriftServ
   override def runQueries(
       queries: Seq[String],
       testCase: TestCase,
-      configSet: Seq[(String, String)]): Unit = {
+      configSet: Seq[(String, String)],
+      testTables: Seq[String]): Unit = {
     // We do not test with configSet.
     withJdbcStatement { statement =>
 
-      loadTestData(statement)
+      loadTestData(statement, testTables)
 
       configSet.foreach { case (k, v) =>
         statement.execute(s"SET $k = $v")
@@ -263,58 +264,70 @@ class ThriftServerQueryTestSuite extends SQLQueryTestSuite with SharedThriftServ
   }
 
   /** Load built-in test tables. */
-  private def loadTestData(statement: Statement): Unit = {
+  private def loadTestData(statement: Statement, testTables: Seq[String]): Unit = {
     // Prepare the data
-    statement.execute(
-      """
-        |CREATE OR REPLACE TEMPORARY VIEW testdata as
-        |SELECT id AS key, CAST(id AS string) AS value FROM range(1, 101)
-      """.stripMargin)
-    statement.execute(
-      """
-        |CREATE OR REPLACE TEMPORARY VIEW arraydata as
-        |SELECT * FROM VALUES
-        |(ARRAY(1, 2, 3), ARRAY(ARRAY(1, 2, 3))),
-        |(ARRAY(2, 3, 4), ARRAY(ARRAY(2, 3, 4))) AS v(arraycol, nestedarraycol)
-      """.stripMargin)
-    statement.execute(
-      """
-        |CREATE OR REPLACE TEMPORARY VIEW mapdata as
-        |SELECT * FROM VALUES
-        |MAP(1, 'a1', 2, 'b1', 3, 'c1', 4, 'd1', 5, 'e1'),
-        |MAP(1, 'a2', 2, 'b2', 3, 'c2', 4, 'd2'),
-        |MAP(1, 'a3', 2, 'b3', 3, 'c3'),
-        |MAP(1, 'a4', 2, 'b4'),
-        |MAP(1, 'a5') AS v(mapcol)
-      """.stripMargin)
-    statement.execute(
-      s"""
-         |CREATE TEMPORARY VIEW aggtest
-         |  (a int, b float)
-         |USING csv
-         |OPTIONS (path '${baseResourcePath.getParent}/test-data/postgresql/agg.data',
-         |  header 'false', delimiter '\t')
-      """.stripMargin)
-    statement.execute(
-      s"""
-         |CREATE OR REPLACE TEMPORARY VIEW onek
-         |  (unique1 int, unique2 int, two int, four int, ten int, twenty int, hundred int,
-         |    thousand int, twothousand int, fivethous int, tenthous int, odd int, even int,
-         |    stringu1 string, stringu2 string, string4 string)
-         |USING csv
-         |OPTIONS (path '${baseResourcePath.getParent}/test-data/postgresql/onek.data',
-         |  header 'false', delimiter '\t')
-      """.stripMargin)
-    statement.execute(
-      s"""
-         |CREATE OR REPLACE TEMPORARY VIEW tenk1
-         |  (unique1 int, unique2 int, two int, four int, ten int, twenty int, hundred int,
-         |    thousand int, twothousand int, fivethous int, tenthous int, odd int, even int,
-         |    stringu1 string, stringu2 string, string4 string)
-         |USING csv
-         |  OPTIONS (path '${baseResourcePath.getParent}/test-data/postgresql/tenk.data',
-         |  header 'false', delimiter '\t')
-      """.stripMargin)
+    if (testTables.contains("testdata")) {
+      statement.execute(
+        """
+          |CREATE OR REPLACE TEMPORARY VIEW testdata as
+          |SELECT id AS key, CAST(id AS string) AS value FROM range(1, 101)
+        """.stripMargin)
+    }
+    if (testTables.contains("arraydata")) {
+      statement.execute(
+        """
+          |CREATE OR REPLACE TEMPORARY VIEW arraydata as
+          |SELECT * FROM VALUES
+          |(ARRAY(1, 2, 3), ARRAY(ARRAY(1, 2, 3))),
+          |(ARRAY(2, 3, 4), ARRAY(ARRAY(2, 3, 4))) AS v(arraycol, nestedarraycol)
+        """.stripMargin)
+    }
+    if (testTables.contains("mapdata")) {
+      statement.execute(
+        """
+          |CREATE OR REPLACE TEMPORARY VIEW mapdata as
+          |SELECT * FROM VALUES
+          |MAP(1, 'a1', 2, 'b1', 3, 'c1', 4, 'd1', 5, 'e1'),
+          |MAP(1, 'a2', 2, 'b2', 3, 'c2', 4, 'd2'),
+          |MAP(1, 'a3', 2, 'b3', 3, 'c3'),
+          |MAP(1, 'a4', 2, 'b4'),
+          |MAP(1, 'a5') AS v(mapcol)
+        """.stripMargin)
+    }
+    if (testTables.contains("aggtest")) {
+      statement.execute(
+        s"""
+           |CREATE TEMPORARY VIEW aggtest
+           |  (a int, b float)
+           |USING csv
+           |OPTIONS (path '${baseResourcePath.getParent}/test-data/postgresql/agg.data',
+           |  header 'false', delimiter '\t')
+        """.stripMargin)
+    }
+    if (testTables.contains("onek")) {
+      statement.execute(
+        s"""
+           |CREATE OR REPLACE TEMPORARY VIEW onek
+           |  (unique1 int, unique2 int, two int, four int, ten int, twenty int, hundred int,
+           |    thousand int, twothousand int, fivethous int, tenthous int, odd int, even int,
+           |    stringu1 string, stringu2 string, string4 string)
+           |USING csv
+           |OPTIONS (path '${baseResourcePath.getParent}/test-data/postgresql/onek.data',
+           |  header 'false', delimiter '\t')
+        """.stripMargin)
+    }
+    if (testTables.contains("tenk1")) {
+      statement.execute(
+        s"""
+           |CREATE OR REPLACE TEMPORARY VIEW tenk1
+           |  (unique1 int, unique2 int, two int, four int, ten int, twenty int, hundred int,
+           |    thousand int, twothousand int, fivethous int, tenthous int, odd int, even int,
+           |    stringu1 string, stringu2 string, string4 string)
+           |USING csv
+           |  OPTIONS (path '${baseResourcePath.getParent}/test-data/postgresql/tenk.data',
+           |  header 'false', delimiter '\t')
+        """.stripMargin)
+    }
   }
 
   // Returns true if sql is retrieving data.