From 6119dfcb4cef60a84912be674800cdd2f28a5602 Mon Sep 17 00:00:00 2001
From: ga-ram <galam.choi@navercorp.com>
Date: Mon, 9 Jan 2023 14:44:26 +0900
Subject: [PATCH] [#9614] Apply temporary fix to apache/pinot#6834 issue

---
 .../WrappedPinotPreparedStatement.java        | 125 +++++++++++++++++-
 1 file changed, 124 insertions(+), 1 deletion(-)

diff --git a/commons-pinot/src/main/java/com/navercorp/pinpoint/common/pinot/datasource/WrappedPinotPreparedStatement.java b/commons-pinot/src/main/java/com/navercorp/pinpoint/common/pinot/datasource/WrappedPinotPreparedStatement.java
index e39316e38d25..6226f2186b68 100644
--- a/commons-pinot/src/main/java/com/navercorp/pinpoint/common/pinot/datasource/WrappedPinotPreparedStatement.java
+++ b/commons-pinot/src/main/java/com/navercorp/pinpoint/common/pinot/datasource/WrappedPinotPreparedStatement.java
@@ -16,30 +16,153 @@
 
 package com.navercorp.pinpoint.common.pinot.datasource;
 
+import org.apache.commons.codec.binary.Hex;
 import org.apache.pinot.client.PinotConnection;
 import org.apache.pinot.client.PinotPreparedStatement;
+import org.apache.pinot.client.utils.DateTimeUtils;
+import org.apache.pinot.client.utils.DriverUtils;
 
+import java.math.BigDecimal;
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Time;
+import java.sql.Timestamp;
+import java.sql.Date;
 
 /**
  * @author Hyunjoon Cho
  */
 public class WrappedPinotPreparedStatement extends PinotPreparedStatement {
+    // Temporary variables for fillStatementWithParameters(). TODO: remove these when Pinot driver is fixed.
+    private String query;
+    private final String[] parameters;
+    private final int maxRows = Integer.MAX_VALUE;
 
     public WrappedPinotPreparedStatement(PinotConnection connection, String query) {
         super(connection, query);
+        this.query = query;
+        if (!DriverUtils.queryContainsLimitStatement(this.query)) {
+            this.query = this.query + " LIMIT " + this.maxRows;
+        }
+        this.parameters = new String[this.getQuestionMarkCount(query)];
+    }
+
+    /* copied from org.apache.pinot.client.PreparedStatement. TODO: remove when Pinot driver is fixed. */
+    private int getQuestionMarkCount(String query) {
+        int questionMarkCount = 0;
+
+        for(int index = query.indexOf(63); index != -1; index = query.indexOf(63, index + 1)) {
+            ++questionMarkCount;
+        }
+
+        return questionMarkCount;
     }
 
+
     @Override
     public boolean execute() throws SQLException {
         // WARNING Do not invoke close().
         // Ignore spotbug warning
-        ResultSet resultSet = executeQuery();
+        ResultSet resultSet = executeQuery(fillStatementWithParameters());
         if (resultSet.isLast()) {
             return false;
         } else {
             return true;
         }
     }
+
+    /* temporary function. TODO: remove these when Pinot driver is fixed */
+    private String fillStatementWithParameters() throws SQLException {
+        String[] queryParts = query.split("\\?");
+        StringBuilder sb = new StringBuilder();
+        int i = 0;
+        while (i < parameters.length) {
+            sb.append(queryParts[i]);
+            if (parameters[i] == null) {
+                throw new SQLException("Prepared Statement parameter is not provided. index = " + i);
+
+            }
+            sb.append(parameters[i]);
+            i++;
+        }
+
+        while (i < queryParts.length) {
+            sb.append(queryParts[i]);
+            i++;
+        }
+
+        return sb.toString();
+    }
+
+    /* temporary setter functions to redirect mybatis parameter setting to temporary variables this.query and this.parameters TODO: remove these when Pinot driver is fixed */
+    public void setNull(int parameterIndex, int sqlType) throws SQLException {
+        this.validateState();
+        this.parameters[parameterIndex - 1] = "'NULL'";
+    }
+
+    public void setBoolean(int parameterIndex, boolean x) throws SQLException {
+        this.validateState();
+        this.parameters[parameterIndex - 1] =  "'" + x + "'";
+    }
+
+    public void setShort(int parameterIndex, short x) throws SQLException {
+        this.validateState();
+        this.parameters[parameterIndex - 1] = String.valueOf(x);
+    }
+
+    public void setInt(int parameterIndex, int x) throws SQLException {
+        this.validateState();
+        this.parameters[parameterIndex - 1] = String.valueOf(x);
+    }
+
+    public void setLong(int parameterIndex, long x) throws SQLException {
+        this.validateState();
+        this.parameters[parameterIndex - 1] = String.valueOf(x);
+    }
+
+    public void setFloat(int parameterIndex, float x) throws SQLException {
+        this.validateState();
+        this.parameters[parameterIndex - 1] = String.valueOf(x);
+    }
+
+    public void setDouble(int parameterIndex, double x) throws SQLException {
+        this.validateState();
+        this.parameters[parameterIndex - 1] = String.valueOf(x);
+    }
+
+    public void setString(int parameterIndex, String x) throws SQLException {
+        this.validateState();
+        this.parameters[parameterIndex - 1] =  "'" + x.replace("'", "''") + "'";;
+    }
+
+    public void setBytes(int parameterIndex, byte[] x) throws SQLException {
+        this.validateState();
+        this.parameters[parameterIndex - 1] = Hex.encodeHexString(x);
+    }
+
+    public void setDate(int parameterIndex, Date x) throws SQLException {
+        this.parameters[parameterIndex - 1] = DateTimeUtils.dateToString(x);
+    }
+
+    public void setTime(int parameterIndex, Time x) throws SQLException {
+        this.validateState();
+        this.parameters[parameterIndex - 1] = DateTimeUtils.timeToString(x);
+    }
+
+    public void setTimestamp(int parameterIndex, Timestamp x) throws SQLException {
+        this.validateState();
+        this.parameters[parameterIndex - 1] = DateTimeUtils.timeStampToString(x);
+    }
+
+    public void setBigDecimal(int parameterIndex, BigDecimal x) throws SQLException {
+        this.validateState();
+        this.parameters[parameterIndex - 1] = x.toString();
+    }
+
+    public void clearParameters() throws SQLException {
+        this.validateState();
+        for (int i = 0; i < this.parameters.length; i++) {
+            this.parameters[i] = null;
+        }
+    }
 }