CriteriaBuilder extension fort tuples IN feature

CriteriaBuilderTupleIn.java

package dev.joserg.jpa;

import java.util.ArrayList;
import java.util.List;
import java.util.function.Function;

import jakarta.persistence.criteria.CriteriaBuilder;
import jakarta.persistence.criteria.Expression;
import jakarta.persistence.criteria.Predicate;
import lombok.extern.slf4j.Slf4j;

@Slf4j
public class CriteriaBuilderTupleIn {

  private final CriteriaBuilder criteriaBuilder;

  public CriteriaBuilderTupleIn(CriteriaBuilder criteriaBuilder) {
    this.criteriaBuilder = criteriaBuilder;
  }

  public <T> Predicate tupleIn(final List<T> tupables, final Function<T, In> tupleizer) {
    return this.tupleIn(tupables.stream().map(tupable -> InTupable.of(tupable, tupleizer)).toList());
  }

  public Predicate tupleIn(final List<? extends InTupable> inTupables) {
    final var tuples = inTupables.stream().map(InTupable::inTupleize).toList();
    return tuples.stream().findFirst()
        .map(tuple -> this.tupleIn(tuple.getExpressions(), tuples.stream().map(In::getValues).toList()))
        .orElseGet(this.criteriaBuilder::disjunction);
  }

  public Predicate tupleIn(final List<? extends Expression<?>> tupleExpression, final List<List<Object>> tupleValues) {
    final var leftIn = this.criteriaBuilder.function("", Object.class, tupleExpression.toArray(Expression[]::new));
    final var rightIn = tupleValues.stream()
        .map(tupleValue -> {
          if (tupleExpression.size() != tupleValue.size()) {
            throw new IllegalArgumentException("Tuple expression and tuple value must have the same size");
          }
          final var values = tupleValue.stream().map(this.criteriaBuilder::literal).toArray(Expression[]::new);
          return this.criteriaBuilder.function("", Object.class, values);
        })
        .toArray(Expression[]::new);
    try {
      return leftIn.in(rightIn);
    } catch (final Exception e) {
      log.warn("Build predicate with multiple columns IN clause failed. Using simulated feature");
      return this.simulatedTupleIn(tupleExpression, tupleValues);
    }
  }

  public <T> Predicate simulatedTupleIn(final List<T> tupables, final Function<T, In> tupleizer) {
    return this.simulatedTupleIn(tupables.stream().map(tupable -> InTupable.of(tupable, tupleizer)).toList());
  }

  public Predicate simulatedTupleIn(final List<? extends InTupable> inTupables) {
    final var tuples = inTupables.stream().map(InTupable::inTupleize).toList();
    return tuples.stream().findFirst()
        .map(tuple -> this.tupleIn(tuple.getExpressions(), tuples.stream().map(In::getValues).toList()))
        .orElseGet(this.criteriaBuilder::disjunction);
  }

  public Predicate simulatedTupleIn(final List<? extends Expression<?>> tupleExpression, final List<List<Object>> tupleValues) {
    var outerPredicate = this.criteriaBuilder.disjunction();
    for (final var tupleValue : tupleValues) {
      if (tupleExpression.size() != tupleValue.size()) {
        throw new IllegalArgumentException("Tuple expression and tuple value must have the same size");
      }
      var innerPredicate = this.criteriaBuilder.conjunction();
      for (int i = 0; i < tupleExpression.size(); i++) {
        final var equalPredicate = this.criteriaBuilder.equal(tupleExpression.get(i), tupleValue.get(i));
        innerPredicate = this.criteriaBuilder.and(innerPredicate, equalPredicate);
      }
      outerPredicate = this.criteriaBuilder.or(outerPredicate, innerPredicate);
    }
    return outerPredicate;
  }

  public interface InTupable {
    static <T> InTupable of(final T tupable, final Function<T, In> tupleizer) {
      return () -> tupleizer.apply(tupable);
    }

    In inTupleize();
  }

  public static class In {

    private final List<Expression<?>> expressions;

    private final List<Object> values;

    public In() {
      this.expressions = new ArrayList<>();
      this.values = new ArrayList<>();
    }

    public <T> In element(Expression<T> expression, T value) {
      this.expressions.add(expression);
      this.values.add(value);
      return this;
    }

    public List<Expression<?>> getExpressions() {
      return this.expressions;
    }

    public List<Object> getValues() {
      return this.values;
    }
  }

}