From e3fdbfb5e20f3d0aec692ac56ba7014d304b8a74 Mon Sep 17 00:00:00 2001 From: linzch3 <20380442+linzch3@users.noreply.github.com> Date: Wed, 4 Nov 2026 10:44:07 +0000 Subject: [PATCH] GH-1439: Release parameter memory when binding fails If binding a parameter value throws, the VectorSchemaRoot holding the values bound so far has not been handed to the prepared statement yet, so nothing closes it. Close it before rethrowing. Generated-by: Grok Bot (AI coding assistant) --- .../jdbc/utils/AvaticaParameterBinder.java | 15 +- .../utils/AvaticaParameterBinderTest.java | 164 ++++++++++++++++++ 2 files changed, 177 insertions(+), 2 deletions(-) create mode 100644 flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinderTest.java diff --git a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java index 8f40d6698e..e2a05c3c44 100644 --- a/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java +++ b/flight/flight-sql-jdbc-core/src/main/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinder.java @@ -61,6 +61,7 @@ public class AvaticaParameterBinder { private final PreparedStatement preparedStatement; private final VectorSchemaRoot parameters; + private boolean parametersHandedOff = false; /** * Instantiate a new AvaticaParameterBinder. @@ -98,13 +99,23 @@ public void bind(List typedValues, int index) { preparedStatement.getParameterSchema().getFields().size(), typedValues.size())); } - for (int i = 0; i < typedValues.size(); i++) { - bind(parameters.getVector(i), typedValues.get(i), index); + try { + for (int i = 0; i < typedValues.size(); i++) { + bind(parameters.getVector(i), typedValues.get(i), index); + } + } catch (RuntimeException e) { + // Until the root is handed to the prepared statement, nothing else owns it, so release the + // memory held by any values bound before the failure. + if (!parametersHandedOff) { + parameters.close(); + } + throw e; } if (!typedValues.isEmpty()) { parameters.setRowCount(index + 1); preparedStatement.setParameters(parameters); + parametersHandedOff = true; } } diff --git a/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinderTest.java b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinderTest.java new file mode 100644 index 0000000000..9375d460ae --- /dev/null +++ b/flight/flight-sql-jdbc-core/src/test/java/org/apache/arrow/driver/jdbc/utils/AvaticaParameterBinderTest.java @@ -0,0 +1,164 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.arrow.driver.jdbc.utils; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.util.Arrays; +import java.util.Collections; +import org.apache.arrow.driver.jdbc.client.ArrowFlightSqlClientHandler; +import org.apache.arrow.flight.FlightInfo; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.VectorSchemaRoot; +import org.apache.arrow.vector.types.pojo.ArrowType; +import org.apache.arrow.vector.types.pojo.Field; +import org.apache.arrow.vector.types.pojo.FieldType; +import org.apache.arrow.vector.types.pojo.Schema; +import org.apache.calcite.avatica.ColumnMetaData; +import org.apache.calcite.avatica.Meta; +import org.apache.calcite.avatica.remote.TypedValue; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Tests for {@link AvaticaParameterBinder}. */ +public class AvaticaParameterBinderTest { + + private static final Schema PARAMETER_SCHEMA = + new Schema( + Arrays.asList( + new Field("text", FieldType.notNullable(new ArrowType.Utf8()), null), + new Field("number", FieldType.notNullable(new ArrowType.Int(64, true)), null))); + + private BufferAllocator allocator; + + @BeforeEach + public void setUp() { + allocator = new RootAllocator(); + } + + @AfterEach + public void tearDown() { + allocator.close(); + } + + @Test + public void testFailedBindReleasesPreviouslyBoundValues() { + for (int i = 0; i < 3; i++) { + try (StubPreparedStatement statement = new StubPreparedStatement()) { + AvaticaParameterBinder binder = new AvaticaParameterBinder(statement, allocator); + UnsupportedOperationException e = + assertThrows( + UnsupportedOperationException.class, + () -> + binder.bind( + Arrays.asList( + TypedValue.ofLocal(ColumnMetaData.Rep.STRING, "O'Reilly"), null))); + assertEquals("Can't set null on non-nullable parameter", e.getMessage()); + } + assertEquals(0, allocator.getAllocatedMemory()); + } + } + + @Test + public void testFailedBindOnTypeMismatchReleasesPreviouslyBoundValues() { + try (StubPreparedStatement statement = new StubPreparedStatement()) { + AvaticaParameterBinder binder = new AvaticaParameterBinder(statement, allocator); + assertThrows( + UnsupportedOperationException.class, + () -> + binder.bind( + Arrays.asList( + TypedValue.ofLocal(ColumnMetaData.Rep.STRING, "text"), + TypedValue.ofLocal(ColumnMetaData.Rep.STRING, "not a number")))); + } + assertEquals(0, allocator.getAllocatedMemory()); + } + + @Test + public void testFailedBatchBindAfterHandOffKeepsStatementParameters() { + try (StubPreparedStatement statement = new StubPreparedStatement()) { + AvaticaParameterBinder binder = new AvaticaParameterBinder(statement, allocator); + binder.bind( + Arrays.asList( + TypedValue.ofLocal(ColumnMetaData.Rep.STRING, "text"), + TypedValue.ofLocal(ColumnMetaData.Rep.LONG, 42L)), + 0); + VectorSchemaRoot handedOff = statement.parameters; + assertThrows( + UnsupportedOperationException.class, + () -> + binder.bind( + Arrays.asList(TypedValue.ofLocal(ColumnMetaData.Rep.STRING, "more"), null), 1)); + // The statement owns the root after the first row, so it must still be usable. + assertSame(handedOff, statement.parameters); + assertEquals(1, statement.parameters.getRowCount()); + } + assertEquals(0, allocator.getAllocatedMemory()); + } + + /** A prepared statement that owns the parameters it is given, like the Flight SQL client. */ + private static final class StubPreparedStatement + implements ArrowFlightSqlClientHandler.PreparedStatement { + private VectorSchemaRoot parameters; + + @Override + public FlightInfo executeQuery() { + throw new UnsupportedOperationException(); + } + + @Override + public long executeUpdate() { + throw new UnsupportedOperationException(); + } + + @Override + public Meta.StatementType getType() { + return Meta.StatementType.SELECT; + } + + @Override + public Schema getDataSetSchema() { + return new Schema(Collections.emptyList()); + } + + @Override + public Schema getParameterSchema() { + return PARAMETER_SCHEMA; + } + + @Override + public Boolean isUpdate() { + return false; + } + + @Override + public void setParameters(VectorSchemaRoot parameters) { + this.parameters = parameters; + } + + @Override + public void close() { + if (parameters != null) { + parameters.close(); + } + } + } +}