Examples with SqlResultConsumer io.micronaut.data.jdbc.mapper.SqlResultConsumer used on opensource projects

Example 1 with SqlResultConsumer

use of io.micronaut.data.jdbc.mapper.SqlResultConsumer in project micronaut-data by micronaut-projects.

the class DefaultJdbcRepositoryOperations method findStream.

private <T, R> Stream<R> findStream(@NonNull PreparedQuery<T, R> preparedQuery, Connection connection) {
    Class<R> resultType = preparedQuery.getResultType();
    AtomicBoolean finished = new AtomicBoolean();
    PreparedStatement ps;
    try {
        ps = prepareStatement(connection, connection::prepareStatement, preparedQuery, false, false);
    } catch (Exception e) {
        throw new DataAccessException("SQL Error preparing Query: " + e.getMessage(), e);
    }
    ResultSet openedRs = null;
    ResultSet rs;
    try {
        openedRs = ps.executeQuery();
        rs = openedRs;
        boolean dtoProjection = preparedQuery.isDtoProjection();
        boolean isEntity = preparedQuery.getResultDataType() == DataType.ENTITY;
        Spliterator<R> spliterator;
        if (isEntity || dtoProjection) {
            SqlResultConsumer sqlMappingConsumer = preparedQuery.hasResultConsumer() ? preparedQuery.getParameterInRole(SqlResultConsumer.ROLE, SqlResultConsumer.class).orElse(null) : null;
            SqlTypeMapper<ResultSet, R> mapper;
            final RuntimePersistentEntity<R> persistentEntity = getEntity(resultType);
            if (dtoProjection) {
                mapper = new SqlDTOMapper<>(persistentEntity, columnNameResultSetReader, jsonCodec, conversionService);
            } else {
                Set<JoinPath> joinFetchPaths = preparedQuery.getJoinFetchPaths();
                SqlResultEntityTypeMapper<ResultSet, R> entityTypeMapper = new SqlResultEntityTypeMapper<>(persistentEntity, columnNameResultSetReader, joinFetchPaths, jsonCodec, (loadedEntity, o) -> {
                    if (loadedEntity.hasPostLoadEventListeners()) {
                        return triggerPostLoad(o, loadedEntity, preparedQuery.getAnnotationMetadata());
                    } else {
                        return o;
                    }
                }, conversionService);
                boolean onlySingleEndedJoins = isOnlySingleEndedJoins(getEntity(preparedQuery.getRootEntity()), joinFetchPaths);
                // Cannot stream ResultSet for "many" joined query
                if (!onlySingleEndedJoins) {
                    try {
                        SqlResultEntityTypeMapper.PushingMapper<ResultSet, List<R>> manyMapper = entityTypeMapper.readAllWithJoins();
                        while (rs.next()) {
                            manyMapper.processRow(rs);
                        }
                        return manyMapper.getResult().stream();
                    } finally {
                        closeResultSet(ps, rs, finished);
                    }
                } else {
                    mapper = entityTypeMapper;
                }
            }
            spliterator = new Spliterators.AbstractSpliterator<R>(Long.MAX_VALUE, Spliterator.ORDERED | Spliterator.IMMUTABLE) {

                @Override
                public boolean tryAdvance(Consumer<? super R> action) {
                    if (finished.get()) {
                        return false;
                    }
                    boolean hasNext = mapper.hasNext(rs);
                    if (hasNext) {
                        R o = mapper.map(rs, resultType);
                        if (sqlMappingConsumer != null) {
                            sqlMappingConsumer.accept(rs, o);
                        }
                        action.accept(o);
                    } else {
                        closeResultSet(ps, rs, finished);
                    }
                    return hasNext;
                }
            };
        } else {
            spliterator = new Spliterators.AbstractSpliterator<R>(Long.MAX_VALUE, Spliterator.ORDERED | Spliterator.IMMUTABLE) {

                @Override
                public boolean tryAdvance(Consumer<? super R> action) {
                    if (finished.get()) {
                        return false;
                    }
                    try {
                        boolean hasNext = rs.next();
                        if (hasNext) {
                            Object v = columnIndexResultSetReader.readDynamic(rs, 1, preparedQuery.getResultDataType());
                            if (resultType.isInstance(v)) {
                                // noinspection unchecked
                                action.accept((R) v);
                            } else if (v != null) {
                                Object r = columnIndexResultSetReader.convertRequired(v, resultType);
                                if (r != null) {
                                    action.accept((R) r);
                                }
                            }
                        } else {
                            closeResultSet(ps, rs, finished);
                        }
                        return hasNext;
                    } catch (SQLException e) {
                        throw new DataAccessException("Error retrieving next JDBC result: " + e.getMessage(), e);
                    }
                }
            };
        }
        return StreamSupport.stream(spliterator, false).onClose(() -> {
            closeResultSet(ps, rs, finished);
        });
    } catch (Exception e) {
        closeResultSet(ps, openedRs, finished);
        throw new DataAccessException("SQL Error executing Query: " + e.getMessage(), e);
    }
}