<commit_before>package com.adgear.anoa.read;

import com.adgear.anoa.Anoa;
import com.adgear.anoa.AnoaHandler;
import com.fasterxml.jackson.core.ObjectCodec;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.fasterxml.jackson.databind.util.TokenBuffer;

import org.apache.avro.Schema;
import org.apache.commons.codec.digest.DigestUtils;
import org.codehaus.jackson.node.JsonNodeFactory;
import org.jooq.lambda.SQL;
import org.jooq.lambda.Unchecked;
import org.jooq.lambda.fi.util.function.CheckedFunction;

import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.Types;
import java.util.List;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;

/**
 * Utility class for streaming Jackson records from a JDBC result set.
 */
public class JdbcStreams {

  /**
   * Create with default object mapper
   */
  public JdbcStreams() {
    this(new ObjectMapper());
  }

  /**
   * @param objectCodec Object mapper to use
   */
  public JdbcStreams(ObjectCodec objectCodec) {
    this.objectCodec = objectCodec;
  }

  /**
   * Object mapper used by this instance
   */
  final public ObjectCodec objectCodec;

  /**
   * @param resultSet the JDBC result set to scan
   * @return A stream of Jackson records which map to the result set rows.
   */
  public /*@NonNull*/ Stream<ObjectNode> resultSet(/*@NonNull*/ ResultSet resultSet) {
    final Function<ResultSet, ObjectNode> fn;
    try {
      fn = Unchecked.function(rowFn(resultSet.getMetaData()));
    } catch (SQLException e) {
      throw new RuntimeException(e);
    }
    return SQL.seq(resultSet, fn);
  }

  /**
   *
   * @param anoaHandler {@code AnoaHandler} instance to use for exception handling
   * @param resultSet the JDBC result set to scan
   * @param <M> Metadata type
   * @return A stream of Jackson records which map to the result set rows.
   */
  public <M> /*@NonNull*/ Stream<Anoa<ObjectNode, M>> resultSet(
      /*@NonNull*/ AnoaHandler<M> anoaHandler,
      /*@NonNull*/ ResultSet resultSet) {
    final Function<Anoa<ResultSet, M>, Anoa<ObjectNode, M>> fn;
    try {
      fn = anoaHandler.functionChecked(rowFn(resultSet.getMetaData()));
    } catch (SQLException e) {
      throw new RuntimeException(e);
    }
    return SQL.seq(resultSet, fn.compose(anoaHandler::ofNullable));
  }

  protected CheckedFunction<ResultSet, ObjectNode> rowFn(
      ResultSetMetaData rsmd) throws SQLException {
    int n = rsmd.getColumnCount();
    String[] names = new String[n];
    for (int c = 0; c < n; c++) {
      names[c] = rsmd.getColumnLabel(c + 1);
    }
    return (ResultSet resultSet) -> {
      TokenBuffer tb = new TokenBuffer(objectCodec, false);
      tb.writeStartObject();
      for (int c = 0; c < n;) {
        tb.writeFieldName(names[c++]);
        tb.writeObject(resultSet.getObject(c));
      }
      tb.writeEndObject();
      return tb.asParser(objectCodec).readValueAsTree();
    };
  }

  /**
   * @param rsmd JDBC Result Set metadata
   * @return equivalent Avro Schema
   * @throws SQLException
   */
  static public /*@NonNull*/ Schema induceSchema(/*@NonNull*/ ResultSetMetaData rsmd) throws SQLException {
    List<Schema.Field> fields = IntStream.range(1, rsmd.getColumnCount() + 1)
        .mapToObj(Unchecked.intFunction(c -> {<commit_msg>default to string in JdbcStreams<commit_after>package com.adgear.anoa.read;

import com.adgear.anoa.Anoa;
import com.adgear.anoa.AnoaHandler;
import com.fasterxml.jackson.core.ObjectCodec;
import com.fasterxml.jackson.core.TreeNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.node.NullNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import com.fasterxml.jackson.databind.util.TokenBuffer;

import org.apache.avro.Schema;
import org.apache.commons.codec.digest.DigestUtils;
import org.codehaus.jackson.node.JsonNodeFactory;
import org.jooq.lambda.SQL;
import org.jooq.lambda.Unchecked;
import org.jooq.lambda.fi.util.function.CheckedFunction;

import java.io.IOException;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.sql.SQLException;
import java.sql.Types;
import java.util.List;
import java.util.function.Function;
import java.util.stream.Collectors;
import java.util.stream.IntStream;
import java.util.stream.Stream;

/**
 * Utility class for streaming Jackson records from a JDBC result set.
 */
public class JdbcStreams {

  /**
   * Create with default object mapper
   */
  public JdbcStreams() {
    this(new ObjectMapper());
  }

  /**
   * @param objectCodec Object mapper to use
   */
  public JdbcStreams(ObjectCodec objectCodec) {
    this.objectCodec = objectCodec;
  }

  /**
   * Object mapper used by this instance
   */
  final public ObjectCodec objectCodec;

  /**
   * @param resultSet the JDBC result set to scan
   * @return A stream of Jackson records which map to the result set rows.
   */
  public /*@NonNull*/ Stream<ObjectNode> resultSet(/*@NonNull*/ ResultSet resultSet) {
    final Function<ResultSet, ObjectNode> fn;
    try {
      fn = Unchecked.function(rowFn(resultSet.getMetaData()));
    } catch (SQLException e) {
      throw new RuntimeException(e);
    }
    return SQL.seq(resultSet, fn);
  }

  /**
   *
   * @param anoaHandler {@code AnoaHandler} instance to use for exception handling
   * @param resultSet the JDBC result set to scan
   * @param <M> Metadata type
   * @return A stream of Jackson records which map to the result set rows.
   */
  public <M> /*@NonNull*/ Stream<Anoa<ObjectNode, M>> resultSet(
      /*@NonNull*/ AnoaHandler<M> anoaHandler,
      /*@NonNull*/ ResultSet resultSet) {
    final Function<Anoa<ResultSet, M>, Anoa<ObjectNode, M>> fn;
    try {
      fn = anoaHandler.functionChecked(rowFn(resultSet.getMetaData()));
    } catch (SQLException e) {
      throw new RuntimeException(e);
    }
    return SQL.seq(resultSet, fn.compose(anoaHandler::ofNullable));
  }

  protected CheckedFunction<ResultSet, ObjectNode> rowFn(
      ResultSetMetaData rsmd) throws SQLException {
    int n = rsmd.getColumnCount();
    String[] names = new String[n];
    for (int c = 0; c < n; c++) {
      names[c] = rsmd.getColumnLabel(c + 1);
    }
    return (ResultSet resultSet) -> {
      TokenBuffer tb = new TokenBuffer(objectCodec, false);
      tb.writeStartObject();
      for (int c = 0; c < n;) {
        tb.writeFieldName(names[c++]);
        tb.writeTree(objectToTree(resultSet.getObject(c)));
      }
      tb.writeEndObject();
      return tb.asParser(objectCodec).readValueAsTree();
    };
  }

  private TreeNode objectToTree(Object object) throws IOException {
    if (object == null) {
      return NullNode.getInstance();
    }
    TokenBuffer tb = new TokenBuffer(objectCodec, false);
    try {
      tb.writeObject(object);
      return tb.asParser().readValueAsTree();
    } catch (IOException e) {
      tb = new TokenBuffer(objectCodec, false);
    }
    tb.writeString(object.toString());
    return tb.asParser().readValueAsTree();
  }

  /**
   * @param rsmd JDBC Result Set metadata
   * @return equivalent Avro Schema
   * @throws SQLException
   */
  static public /*@NonNull*/ Schema induceSchema(/*@NonNull*/ ResultSetMetaData rsmd) throws SQLException {
    List<Schema.Field> fields = IntStream.range(1, rsmd.getColumnCount() + 1)
        .mapToObj(Unchecked.intFunction(c -> {
{'fraction_non_alphanumeric': 0.1195840554592721, 'fraction_numerical': 0.0013331555792560992, 'mean_word_length': 3.5632603406326036, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 31, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 10}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4895573', 'n_tokens_mistral': 2399, 'n_tokens_neox': 2365, 'n_words': 576}