From 1c279300bc884331aef3349d8af18ca0b09da465 Mon Sep 17 00:00:00 2001 From: Odin-BN Date: Fri, 30 May 2025 10:09:46 -0700 Subject: [PATCH] GH-734: Improve CompositeJdbcConsumer error handling and add unit tests --- .../jdbc/consumer/CompositeJdbcConsumer.java | 21 ++++- .../consumer/CompositeJdbcConsumerTest.java | 94 +++++++++++++++++++ 2 files changed, 112 insertions(+), 3 deletions(-) create mode 100644 adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumerTest.java diff --git a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumer.java b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumer.java index 2366116fd..d21f94d77 100644 --- a/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumer.java +++ b/adapter/jdbc/src/main/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumer.java @@ -44,9 +44,24 @@ public void consume(ResultSet rs) throws SQLException, IOException { } catch (Exception e) { if (consumers[i] instanceof BaseConsumer) { BaseConsumer consumer = (BaseConsumer) consumers[i]; - JdbcFieldInfo fieldInfo = - new JdbcFieldInfo(rs.getMetaData(), consumer.columnIndexInResultSet); - ArrowType arrowType = consumer.vector.getMinorType().getType(); + JdbcFieldInfo fieldInfo = null; + ArrowType arrowType = null; + try { + if (rs != null) { + fieldInfo = new JdbcFieldInfo(rs.getMetaData(), consumer.columnIndexInResultSet); + } + } catch (Exception metaEx) { + // doesn't do anything if ResultSet is null + // to return the JdbcConsumerException + } + try { + if (consumer.vector.getMinorType() != null){ + arrowType = consumer.vector.getMinorType().getType(); + } + } catch (Exception typeEx) { + // doesn't do anything if there is an error when getting null with getMinorType() + // to return the JdbcConsumerException + } throw new JdbcConsumerException( "Exception while consuming JDBC value", e, fieldInfo, arrowType); } else { diff --git a/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumerTest.java b/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumerTest.java new file mode 100644 index 000000000..8e16caf8d --- /dev/null +++ b/adapter/jdbc/src/test/java/org/apache/arrow/adapter/jdbc/consumer/CompositeJdbcConsumerTest.java @@ -0,0 +1,94 @@ +/* + * 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.adapter.jdbc.consumer; + +import org.apache.arrow.adapter.jdbc.consumer.exceptions.JdbcConsumerException; +import org.apache.arrow.memory.BufferAllocator; +import org.apache.arrow.memory.RootAllocator; +import org.apache.arrow.vector.IntVector; +import org.apache.arrow.vector.ValueVector; +import java.sql.ResultSet; +import java.sql.SQLException; + +import static org.junit.jupiter.api.Assertions.*; + +import org.junit.jupiter.api.Test; + +public class CompositeJdbcConsumerTest { + // Faulty consumer that simulates a runtime exception during consume() + // This is used to test that the CompositeJdbcConsumer wraps it properly + static class FaultyConsumer extends BaseConsumer { + public FaultyConsumer(ValueVector vector, int index) { + super(vector, index); + } + + @Override + public void consume(ResultSet rs) throws SQLException { + throw new NullPointerException("Simulating NPE"); + } + } + + @Test + public void testHandlesJdbcConsumerExceptionGracefully() throws SQLException { + // Setup: create an IntVector to simulate a consumer with a valid vector + BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); + IntVector intVector = new IntVector("int", allocator); + intVector.allocateNew(); + + // Simulate a failing consumer with valid vector (to test arrowType extraction) + JdbcConsumer mockConsumer = new FaultyConsumer(intVector, 1); + CompositeJdbcConsumer composite = new CompositeJdbcConsumer(new JdbcConsumer[]{mockConsumer}); + + // Use null ResultSet to simulate failure scenario + ResultSet dummyRs = null; + + // Verify: the failure is caught and wrapped in JdbcConsumerException + JdbcConsumerException thrownEx = assertThrows(JdbcConsumerException.class, () -> { + composite.consume(dummyRs); + }); + assertTrue(thrownEx.getMessage().contains("Exception while consuming JDBC value")); + assertNull(thrownEx.getFieldInfo()); + assertNotNull(thrownEx.getArrowType()); // Should be non-null since vector was valid + } + + // Faulty consumer that has a null ValueVector (to test arrowType = null) + public static class FaultyConsumerWIthNullVector extends BaseConsumer { + public FaultyConsumerWIthNullVector(int index) { + super(null, index); + } + + @Override + public void consume(ResultSet rs) throws SQLException { + throw new NullPointerException("Simulating NPE"); + } + } + + @Test + public void testJdbcConsumerExceptionWhenArrowTypeIsNull() throws SQLException { + // Setup: consumer with null vector + JdbcConsumer mockConsumer = new FaultyConsumerWIthNullVector(1); + CompositeJdbcConsumer composite = new CompositeJdbcConsumer(new JdbcConsumer[]{mockConsumer}); + ResultSet dummyRs = null; + + // Verify: when the consumer fails and the vector is null, + // arrowType in JdbcConsumerException should also be null + JdbcConsumerException thrownEx = assertThrows(JdbcConsumerException.class, () -> { + composite.consume(dummyRs); + }); + assertNull(thrownEx.getArrowType()); + } +}