001package org.itsallcode.jdbc.resultset;
002
003import java.sql.ResultSet;
004import java.sql.SQLException;
005import java.util.*;
006import java.util.stream.Stream;
007import java.util.stream.StreamSupport;
008
009import org.itsallcode.jdbc.Context;
010import org.itsallcode.jdbc.UncheckedSQLException;
011import org.itsallcode.jdbc.resultset.generic.SimpleMetaData;
012
013/**
014 * This class wraps a {@link ResultSet} and allows easy iteration via
015 * {@link Iterator}, {@link List} or {@link Stream}.
016 * 
017 * @param <T> row type
018 */
019public class SimpleResultSet<T> implements AutoCloseable, Iterable<T> {
020    private final ResultSet resultSet;
021    private final ContextRowMapper<T> rowMapper;
022    private final Context context;
023    private final AutoCloseable statement;
024    private Iterator<T> iterator;
025
026    /**
027     * Create a new instance.
028     * 
029     * @param context   database context
030     * @param resultSet the underlying result set
031     * @param rowMapper a row mapper for converting each row
032     * @param statement the statement that created the result set. This will be
033     *                  closed when the result set is closed.
034     */
035    public SimpleResultSet(final Context context, final ResultSet resultSet, final ContextRowMapper<T> rowMapper,
036            final AutoCloseable statement) {
037        this.context = context;
038        this.resultSet = resultSet;
039        this.rowMapper = rowMapper;
040        this.statement = statement;
041    }
042
043    /**
044     * Get result set metadata.
045     * 
046     * @return metadata
047     */
048    public SimpleMetaData getMetaData() {
049        return SimpleMetaData.create(this.resultSet);
050    }
051
052    /**
053     * Get in {@link Iterator} of all rows.
054     * 
055     * @return an iterator with all rows.
056     */
057    @Override
058    public Iterator<T> iterator() {
059        if (iterator != null) {
060            throw new IllegalStateException("Only one iterator allowed per ResultSet");
061        }
062        iterator = ResultSetIterator.create(context, this, rowMapper);
063        return iterator;
064    }
065
066    /**
067     * Collect all rows to a list.
068     * 
069     * @return a list with all rows.
070     */
071    public List<T> toList() {
072        try (Stream<T> stream = stream()) {
073            return stream.toList();
074        }
075    }
076
077    /**
078     * Get a stream of all rows.
079     * 
080     * @return a stream with all rows.
081     */
082    public Stream<T> stream() {
083        final Spliterator<T> spliterator = Spliterators.spliteratorUnknownSize(this.iterator(), Spliterator.ORDERED);
084        return StreamSupport.stream(spliterator, false)
085                .onClose(this::close);
086    }
087
088    /**
089     * Close the underlying {@link ResultSet} and the statement that created it.
090     * 
091     * @throws UncheckedSQLException if closing fails.
092     */
093    @Override
094    public void close() {
095        try {
096            resultSet.close();
097        } catch (final SQLException e) {
098            throw new UncheckedSQLException("Error closing resultset", e);
099        }
100        try {
101            statement.close();
102        } catch (final Exception e) {
103            throw new IllegalStateException("Error closing statement: " + e.getMessage(), e);
104        }
105    }
106
107    private boolean next() {
108        try {
109            return resultSet.next();
110        } catch (final SQLException e) {
111            throw new UncheckedSQLException("Error getting next row", e);
112        }
113    }
114
115    private static final class ResultSetIterator<T> implements Iterator<T> {
116        private final Context context;
117        private final SimpleResultSet<T> resultSet;
118        private final ContextRowMapper<T> rowMapper;
119        private boolean hasNext;
120        private int currentRowIndex;
121
122        private ResultSetIterator(final Context context, final SimpleResultSet<T> simpleResultSet,
123                final ContextRowMapper<T> rowMapper,
124                final boolean hasNext) {
125            this.context = context;
126            this.resultSet = simpleResultSet;
127            this.rowMapper = rowMapper;
128            this.hasNext = hasNext;
129        }
130
131        private static <T> Iterator<T> create(final Context context, final SimpleResultSet<T> simpleResultSet,
132                final ContextRowMapper<T> rowMapper) {
133            final boolean firstRowExists = simpleResultSet.next();
134            return new ResultSetIterator<>(context, simpleResultSet, rowMapper, firstRowExists);
135        }
136
137        @Override
138        public boolean hasNext() {
139            return hasNext;
140        }
141
142        @Override
143        public T next() {
144            if (!hasNext) {
145                throw new NoSuchElementException();
146            }
147            final T row = mapRow();
148            hasNext = resultSet.next();
149            currentRowIndex++;
150            return row;
151        }
152
153        private T mapRow() {
154            try {
155                return rowMapper.mapRow(context, resultSet.resultSet, currentRowIndex);
156            } catch (final SQLException e) {
157                throw new UncheckedSQLException("Error mapping row " + currentRowIndex, e);
158            } catch (final RuntimeException e) {
159                throw new IllegalStateException("Error mapping row " + currentRowIndex + ": " + e.getMessage(), e);
160            }
161        }
162    }
163}