Api#2880: Close the underlying iterator in ClosingIterator in hasNext() call (#2881)
diff --git a/api/src/main/java/org/apache/iceberg/io/ClosingIterator.java b/api/src/main/java/org/apache/iceberg/io/ClosingIterator.java
index 26575fe..58fc931 100644
--- a/api/src/main/java/org/apache/iceberg/io/ClosingIterator.java
+++ b/api/src/main/java/org/apache/iceberg/io/ClosingIterator.java
@@ -23,35 +23,38 @@
import java.io.UncheckedIOException;
import java.util.Iterator;
+/**
+ * A convenience wrapper around {@link CloseableIterator}, providing auto-close
+ * functionality when all of the elements in the iterator are consumed.
+ */
public class ClosingIterator<T> implements Iterator<T> {
private final CloseableIterator<T> iterator;
- private boolean shouldClose = false;
+ private boolean isClosed;
public ClosingIterator(CloseableIterator<T> iterator) {
this.iterator = iterator;
-
}
@Override
public boolean hasNext() {
boolean hasNext = iterator.hasNext();
- this.shouldClose = !hasNext;
+ if (!hasNext && !isClosed) {
+ close();
+ }
return hasNext;
}
@Override
public T next() {
- T next = iterator.next();
+ return iterator.next();
+ }
- if (shouldClose) {
- // this will only be called once because iterator.next would throw NoSuchElementException
- try {
- iterator.close();
- } catch (IOException e) {
- throw new UncheckedIOException(e);
- }
+ private void close() {
+ try {
+ iterator.close();
+ isClosed = true;
+ } catch (IOException e) {
+ throw new UncheckedIOException(e);
}
-
- return next;
}
}
diff --git a/api/src/test/java/org/apache/iceberg/io/TestClosingIterator.java b/api/src/test/java/org/apache/iceberg/io/TestClosingIterator.java
new file mode 100644
index 0000000..5aebc2f
--- /dev/null
+++ b/api/src/test/java/org/apache/iceberg/io/TestClosingIterator.java
@@ -0,0 +1,71 @@
+/*
+ * 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.iceberg.io;
+
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+public class TestClosingIterator {
+ @Test
+ public void testEmptyIterator() {
+ CloseableIterator<String> underlying = mock(CloseableIterator.class);
+ ClosingIterator<String> closingIterator = new ClosingIterator<>(underlying);
+ assertFalse(closingIterator.hasNext());
+ }
+
+ @Test
+ public void testHasNextAndNext() {
+ CloseableIterator<String> underlying = mock(CloseableIterator.class);
+ when(underlying.hasNext()).thenReturn(true);
+ when(underlying.next()).thenReturn("hello");
+ ClosingIterator<String> closingIterator = new ClosingIterator<>(underlying);
+ assertTrue(closingIterator.hasNext());
+ assertEquals("hello", closingIterator.next());
+ }
+
+ @Test
+ public void testUnderlyingIteratorCloseWhenElementsAreExhausted() throws Exception {
+ CloseableIterator<String> underlying = mock(CloseableIterator.class);
+ when(underlying.hasNext()).thenReturn(true).thenReturn(false);
+ when(underlying.next()).thenReturn("hello");
+ ClosingIterator<String> closingIterator = new ClosingIterator<>(underlying);
+ assertTrue(closingIterator.hasNext());
+ assertEquals("hello", closingIterator.next());
+
+ assertFalse(closingIterator.hasNext());
+ verify(underlying, times(1)).close();
+ }
+
+ @Test
+ public void testCloseCalledOnceForMultipleHasNextCalls() throws Exception {
+ CloseableIterator<String> underlying = mock(CloseableIterator.class);
+ ClosingIterator<String> closingIterator = new ClosingIterator<>(underlying);
+ assertFalse(closingIterator.hasNext());
+ assertFalse(closingIterator.hasNext());
+ verify(underlying, times(1)).close();
+ }
+}
