add builder pattern

Author: xborder

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightMetaImpl.java

@@ -102,15 +102,13 @@ ArrowFlightPreparedStatement createPreparedStatement(
       final int resultSetConcurrency,
       final int resultSetHoldability)
       throws SQLException {
-    final StatementHandle handle = super.createStatement(connection.handle);
-    return ArrowFlightPreparedStatement.createPrepared(
-        (ArrowFlightConnection) connection,
-        handle,
-        //        null,
-        query,
-        resultSetType,
-        resultSetConcurrency,
-        resultSetHoldability);
+    return ArrowFlightPreparedStatement.builder((ArrowFlightConnection) connection)
+        .withQuery(query)
+        .withGeneratedHandle()
+        .withResultSetType(resultSetType)
+        .withResultSetConcurrency(resultSetConcurrency)
+        .withResultSetHoldability(resultSetHoldability)
+        .build();
   }
 
   @Override
@@ -153,14 +151,14 @@ public ExecuteResult prepareAndExecute(
           && !(statement instanceof ArrowFlightPreparedStatement)) {
         throw new IllegalStateException("Prepared statement not found: " + handle);
       }
+      if (statement instanceof ArrowFlightPreparedStatement) {
+        ((ArrowFlightPreparedStatement) statement).closePreparedResources();
+      }
       final ArrowFlightPreparedStatement preparedStatement =
-          ArrowFlightPreparedStatement.createPrepared(
-              (ArrowFlightConnection) connection,
-              handle,
-              query,
-              statement.getResultSetType(),
-              statement.getResultSetConcurrency(),
-              statement.getResultSetHoldability());
+          ArrowFlightPreparedStatement.builder((ArrowFlightConnection) connection)
+              .withQuery(query)
+              .withExistingStatement(statement)
+              .build();
       return preparedStatement.prepareAndExecute(callback);
     } catch (SQLTimeoutException e) {
       // So far AvaticaStatement(executeInternal) only handles NoSuchStatement and

flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatement.java

@@ -25,6 +25,7 @@
 import org.apache.arrow.util.Preconditions;
 import org.apache.arrow.vector.types.pojo.Schema;
 import org.apache.calcite.avatica.AvaticaPreparedStatement;
+import org.apache.calcite.avatica.AvaticaStatement;
 import org.apache.calcite.avatica.Meta.ExecuteBatchResult;
 import org.apache.calcite.avatica.Meta.ExecuteResult;
 import org.apache.calcite.avatica.Meta.MetaResultSet;
@@ -44,39 +45,19 @@ private ArrowFlightPreparedStatement(
       final ArrowFlightConnection connection,
       final StatementHandle handle,
       final Signature signature,
+      final ArrowFlightSqlClientHandler.PreparedStatement preparedStatement,
       final int resultSetType,
       final int resultSetConcurrency,
       final int resultSetHoldability)
       throws SQLException {
     super(connection, handle, signature, resultSetType, resultSetConcurrency, resultSetHoldability);
+    this.preparedStatement = Preconditions.checkNotNull(preparedStatement);
+    this.handle.signature = signature;
+    setSignature(signature);
   }
 
-  static ArrowFlightPreparedStatement createPrepared(
-      final ArrowFlightConnection connection,
-      final StatementHandle statementHandle,
-      final String query,
-      final int resultSetType,
-      final int resultSetConcurrency,
-      final int resultSetHoldability)
-      throws SQLException {
-    final ArrowFlightSqlClientHandler.PreparedStatement preparedStatement =
-        connection.getClientHandler().prepare(query);
-    final Signature signature =
-        ArrowFlightMetaImpl.buildSignature(
-            query, preparedStatement.getDataSetSchema(), preparedStatement.getParameterSchema());
-    statementHandle.signature = signature;
-
-    final ArrowFlightPreparedStatement statement =
-        new ArrowFlightPreparedStatement(
-            connection,
-            statementHandle,
-            signature,
-            resultSetType,
-            resultSetConcurrency,
-            resultSetHoldability);
-    statement.preparedStatement = Preconditions.checkNotNull(preparedStatement);
-    statement.setSignature(signature);
-    return statement;
+  static Builder builder(final ArrowFlightConnection connection) {
+    return new Builder(connection);
   }
 
   @Override
@@ -173,4 +154,79 @@ private void ensurePrepared() {
       throw new IllegalStateException("PreparedStatement is already closed.");
     }
   }
+
+  static final class Builder {
+    private final ArrowFlightConnection connection;
+    private StatementHandle handle;
+    private String query;
+    private Integer resultSetType;
+    private Integer resultSetConcurrency;
+    private Integer resultSetHoldability;
+    private boolean generateHandle;
+
+    private Builder(final ArrowFlightConnection connection) {
+      this.connection = Preconditions.checkNotNull(connection);
+    }
+
+    Builder withQuery(final String query) {
+      this.query = Preconditions.checkNotNull(query);
+      return this;
+    }
+
+    Builder withGeneratedHandle() {
+      this.generateHandle = true;
+      this.handle = null;
+      return this;
+    }
+
+    Builder withExistingStatement(final AvaticaStatement statement) throws SQLException {
+      Preconditions.checkNotNull(statement);
+      this.generateHandle = false;
+      this.handle = Preconditions.checkNotNull(statement.handle);
+      this.resultSetType = statement.getResultSetType();
+      this.resultSetConcurrency = statement.getResultSetConcurrency();
+      this.resultSetHoldability = statement.getResultSetHoldability();
+      return this;
+    }
+
+    Builder withResultSetType(final int resultSetType) {
+      this.resultSetType = resultSetType;
+      return this;
+    }
+
+    Builder withResultSetConcurrency(final int resultSetConcurrency) {
+      this.resultSetConcurrency = resultSetConcurrency;
+      return this;
+    }
+
+    Builder withResultSetHoldability(final int resultSetHoldability) {
+      this.resultSetHoldability = resultSetHoldability;
+      return this;
+    }
+
+    ArrowFlightPreparedStatement build() throws SQLException {
+      Preconditions.checkNotNull(query);
+      Preconditions.checkNotNull(resultSetType);
+      Preconditions.checkNotNull(resultSetConcurrency);
+      Preconditions.checkNotNull(resultSetHoldability);
+      if (!generateHandle && handle == null) {
+        throw new IllegalStateException("PreparedStatement builder requires a handle.");
+      }
+
+      final ArrowFlightSqlClientHandler.PreparedStatement preparedStatement =
+          connection.getClientHandler().prepare(query);
+      final Signature signature =
+          ArrowFlightMetaImpl.buildSignature(
+              query, preparedStatement.getDataSetSchema(), preparedStatement.getParameterSchema());
+
+      return new ArrowFlightPreparedStatement(
+          connection,
+          generateHandle ? null : handle,
+          signature,
+          preparedStatement,
+          resultSetType,
+          resultSetConcurrency,
+          resultSetHoldability);
+    }
+  }
 }

flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightPreparedStatementTest.java

@@ -21,6 +21,8 @@
 import static org.junit.jupiter.api.Assertions.assertAll;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.nio.charset.StandardCharsets;
@@ -98,6 +100,27 @@ public void testSimpleQueryNoParameterBindingWithExecute() throws SQLException {
     }
   }
 
+  @Test
+  public void testPrepareStatementRegistersCreatedStatementByGeneratedHandle() throws SQLException {
+    final String query = CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD;
+    final ArrowFlightConnection flightConnection = (ArrowFlightConnection) connection;
+
+    try (final PreparedStatement preparedStatement = connection.prepareStatement(query)) {
+      final ArrowFlightPreparedStatement arrowPreparedStatement =
+          (ArrowFlightPreparedStatement) preparedStatement;
+
+      assertNotNull(
+          flightConnection
+              .getMeta()
+              .getPreparedStatementInstanceOrNull(arrowPreparedStatement.handle));
+      assertSame(
+          arrowPreparedStatement,
+          flightConnection
+              .getMeta()
+              .getPreparedStatementInstanceOrNull(arrowPreparedStatement.handle));
+    }
+  }
+
   @Test
   public void testQueryWithParameterBinding() throws SQLException {
     final String query = "Fake query with parameters";

flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/ArrowFlightStatementExecuteTest.java

@@ -22,6 +22,8 @@
 import static org.hamcrest.CoreMatchers.not;
 import static org.hamcrest.CoreMatchers.nullValue;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
 
 import java.sql.Connection;
 import java.sql.ResultSet;
@@ -137,6 +139,21 @@ public void testExecuteShouldRunSelectQuery() throws SQLException {
         is(allOf(equalTo(statement.getLargeUpdateCount()), equalTo(-1L))));
   }
 
+  @Test
+  public void testExecuteReplacesStatementMapEntryWithPreparedStatement() throws SQLException {
+    final ArrowFlightStatement arrowStatement = (ArrowFlightStatement) statement;
+    final ArrowFlightConnection arrowConnection = (ArrowFlightConnection) connection;
+
+    assertThat(statement.execute(SAMPLE_QUERY_CMD), is(true));
+
+    final ArrowFlightPreparedStatement preparedStatement =
+        arrowConnection.getMeta().getPreparedStatementInstanceOrNull(arrowStatement.handle);
+
+    assertNotNull(preparedStatement);
+    assertSame(preparedStatement, arrowConnection.statementMap.get(arrowStatement.handle.id));
+    assertThat(preparedStatement.handle.id, is(equalTo(arrowStatement.handle.id)));
+  }
+
   @Test
   public void testExecuteShouldRunUpdateQueryForSmallUpdate() throws SQLException {
     assertThat(statement.execute(SAMPLE_UPDATE_QUERY), is(false)); // Means this is an UPDATE query.