io.stargate.db.query.RowsImpacted Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of persistence-api Show documentation
There is a newer version: 2.1.0-BETA-19
package io.stargate.db.query;

import com.datastax.oss.driver.shaded.guava.common.base.Preconditions;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import javax.annotation.Nullable;

/**
 * Specify which rows are updated by a given DML query.
 *
 * There is effectively two implementation of this class, corresponding to the 2 ways rows can be
 * updated in CQL: either a set of specific rows (implemented by {@link Keys}), or a range of rows
 * (implemented by {@link Ranges}). Note that in practice, ranges are allowed by DELETE in CQL so
 * far, and INSERT can only insert a single row at a time.
 */
public abstract class RowsImpacted {
  private @Nullable Set partitionKeys; // Lazily Computed

  private RowsImpacted() {}

  public abstract boolean isRanges();

  public boolean isKeys() {
    return !isRanges();
  }

  public Keys asKeys() {
    Preconditions.checkState(isKeys());
    return (Keys) this;
  }

  public Ranges asRanges() {
    Preconditions.checkState(isRanges());
    return (Ranges) this;
  }

  public Set partitionKeys() {
    if (partitionKeys == null) {
      partitionKeys = computePartitionKeys();
    }
    return partitionKeys;
  }

  protected abstract Set computePartitionKeys();

  public static final class Keys extends RowsImpacted {
    private final List primaryKeys;

    public Keys(List primaryKeys) {
      this.primaryKeys = primaryKeys;
    }

    public List primaryKeys() {
      return primaryKeys;
    }

    @Override
    public boolean isRanges() {
      return false;
    }

    @Override
    protected Set computePartitionKeys() {
      // Minor optimization, but that's a common case in practice.
      if (primaryKeys.size() == 1) {
        return Collections.singleton(primaryKeys.get(0).partitionKey());
      }

      Set partitionKeys = new HashSet<>();
      primaryKeys.forEach(pk -> partitionKeys.add(pk.partitionKey()));
      return partitionKeys;
    }

    @Override
    public boolean equals(Object o) {
      if (this == o) {
        return true;
      }
      if (!(o instanceof Keys)) {
        return false;
      }
      Keys keys = (Keys) o;
      return primaryKeys.equals(keys.primaryKeys);
    }

    @Override
    public int hashCode() {
      return Objects.hash(primaryKeys);
    }

    @Override
    public String toString() {
      return primaryKeys.toString();
    }
  }

  public static final class Ranges extends RowsImpacted {
    private final List rowsRanges;

    public Ranges(List rowsRanges) {
      this.rowsRanges = rowsRanges;
    }

    public List ranges() {
      return rowsRanges;
    }

    @Override
    public boolean isRanges() {
      return true;
    }

    @Override
    protected Set computePartitionKeys() {
      Set partitionKeys = new HashSet<>();
      rowsRanges.forEach(pk -> partitionKeys.add(pk.partitionKey()));
      return partitionKeys;
    }

    @Override
    public boolean equals(Object o) {
      if (this == o) {
        return true;
      }
      if (!(o instanceof Ranges)) {
        return false;
      }
      Ranges ranges = (Ranges) o;
      return rowsRanges.equals(ranges.rowsRanges);
    }

    @Override
    public int hashCode() {
      return Objects.hash(rowsRanges);
    }
  }
}