From c9acdc6f8ff5021d5b0b39b278bcfc4e3364e7d5 Mon Sep 17 00:00:00 2001
From: Stamatis Zampetakis <zabetak@gmail.com>
Date: Mon, 26 Feb 2024 14:00:50 +0100
Subject: [PATCH] Test function consistency by comparing Calcite against
 reference values

Add option to bypass Calcite and execute statement using the raw ResultSet.
---
 .../test/DBFunctionConsistencyTest.java       | 72 +++++++++++--------
 1 file changed, 43 insertions(+), 29 deletions(-)

diff --git a/core/src/test/java/org/apache/calcite/test/DBFunctionConsistencyTest.java b/core/src/test/java/org/apache/calcite/test/DBFunctionConsistencyTest.java
index 50559dd3dfc2..30bcf2959103 100644
--- a/core/src/test/java/org/apache/calcite/test/DBFunctionConsistencyTest.java
+++ b/core/src/test/java/org/apache/calcite/test/DBFunctionConsistencyTest.java
@@ -37,6 +37,10 @@
 import java.util.Map;
 
 public class DBFunctionConsistencyTest {
+  /**
+   * Whether to execute a statement using a raw JDBC connection or via Calcite.
+   */
+  private static final Boolean USE_RAW_JDBC_CONNECTION = false;
 
   private enum Type {
     ORACLE {
@@ -74,46 +78,56 @@ static void teardown() {
 
   @ParameterizedTest
   @CsvSource({
-      "SQRT(4),ORACLE",
-      "SQRT(4),POSTGRES_9_6",
-      "SQRT(4),POSTGRES_12_2",
-      "SQRT(4),MYSQL",
-      "SQRT(-1),ORACLE",
-      "SQRT(-1),POSTGRES_9_6",
-      "SQRT(-1),POSTGRES_12_2",
-      "SQRT(-1),MYSQL"
+      "SQRT(4.0),ORACLE,2",
+      "SQRT(4.0),POSTGRES_9_6,2.000000000000000",
+      "SQRT(4.0),POSTGRES_12_2,2.000000000000000",
+      "SQRT(4.0),MYSQL,2",
+      "SQRT(4),ORACLE,2",
+      "SQRT(4),POSTGRES_9_6,2",
+      "SQRT(4),POSTGRES_12_2,2",
+      "SQRT(4),MYSQL,2",
+      "SQRT(-1),ORACLE,ERROR",
+      "SQRT(-1),POSTGRES_9_6,ERROR",
+      "SQRT(-1),POSTGRES_12_2,ERROR",
+      "SQRT(-1),MYSQL,"
   })
-  void testFunction(String function, String db) {
-    String calciteResult = execute(Type.valueOf(db), true, function);
-    String rawResult = execute(Type.valueOf(db), false, function);
-    Assertions.assertEquals(rawResult, calciteResult);
+  void testFunction(String function, String db, String expectedResult) {
+    Assertions.assertEquals(expectedResult, execute(Type.valueOf(db), function));
   }
 
-  String execute(Type dbType, boolean useCalcite, String exp) {
-    JdbcDatabaseContainer<?> db = DBS.get(dbType);
-    final String url;
-    final String query;
-    if (useCalcite) {
-      url = String.format(
-          "jdbc:calcite:schemaType=JDBC; schema.jdbcUser=%s; schema.jdbcPassword=%s; schema"
-              + ".jdbcUrl=%s", db.getUsername(), db.getPassword(), db.getJdbcUrl());
-      query = "SELECT " + exp;
-    } else {
-      url = db.getJdbcUrl();
-      query = dbType.query(exp);
-    }
-    try (Connection c = DriverManager.getConnection(url, db.getUsername(), db.getPassword())) {
-      try (PreparedStatement stmt = c.prepareStatement(query)) {
+  private static String execute(Type dbType, String exp) {
+    try (Connection c = getConnection(dbType)) {
+      try (PreparedStatement stmt = c.prepareStatement(query(dbType, exp))) {
         try (ResultSet rs = stmt.executeQuery()) {
           if (rs.next()) {
             return rs.getString(1);
           } else {
-            return "";
+            throw new AssertionError("NoResult");
           }
         }
       }
     } catch (Exception e) {
-      return e.getMessage();
+      return "ERROR";
+    }
+  }
+
+  private static Connection getConnection(Type dbType) throws SQLException {
+    JdbcDatabaseContainer<?> db = DBS.get(dbType);
+    if (USE_RAW_JDBC_CONNECTION) {
+      return DriverManager.getConnection(db.getJdbcUrl(), db.getUsername(), db.getPassword());
+    } else {
+      String calciteUrl = String.format(
+          "jdbc:calcite:schemaType=JDBC; schema.jdbcUser=%s; schema.jdbcPassword=%s; schema"
+              + ".jdbcUrl=%s", db.getUsername(), db.getPassword(), db.getJdbcUrl());
+      return DriverManager.getConnection(calciteUrl);
+    }
+  }
+
+  private static String query(Type dbType, String exp) {
+    if (USE_RAW_JDBC_CONNECTION) {
+      return dbType.query(exp);
+    } else {
+      return "SELECT " + exp;
     }
   }
 }