All Downloads are FREE. Search and download functionalities are using the official Maven repository.

com.eventsourcing.postgresql.index.PostgreSQLAttributeIndex Maven / Gradle / Ivy

Go to download 

/**
 * Copyright (c) 2016, All Contributors (see CONTRIBUTORS file)
 *
 * This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/.
 */
package com.eventsourcing.postgresql.index;

import com.eventsourcing.Entity;
import com.eventsourcing.EntityHandle;
import com.eventsourcing.ResolvedEntityHandle;
import com.eventsourcing.index.AbstractAttributeIndex;
import com.eventsourcing.index.Attribute;
import com.eventsourcing.index.KeyObjectStore;
import com.eventsourcing.index.MultiValueAttribute;
import com.eventsourcing.layout.Layout;
import com.eventsourcing.layout.SerializableComparable;
import com.eventsourcing.layout.TypeHandler;
import com.eventsourcing.postgresql.PostgreSQLSerialization;
import com.eventsourcing.postgresql.PostgreSQLStatementIterator;
import com.googlecode.cqengine.index.Index;
import com.googlecode.cqengine.index.support.*;
import com.googlecode.cqengine.index.unique.UniqueIndex;
import com.googlecode.cqengine.persistence.support.ObjectSet;
import com.googlecode.cqengine.persistence.support.ObjectStore;
import com.googlecode.cqengine.query.Query;
import com.googlecode.cqengine.query.option.QueryOptions;
import com.googlecode.cqengine.query.simple.Equal;
import com.googlecode.cqengine.query.simple.Has;
import com.googlecode.cqengine.resultset.ResultSet;
import com.googlecode.cqengine.resultset.closeable.CloseableResultSet;
import com.impossibl.postgres.jdbc.PGSQLIntegrityConstraintViolationException;
import lombok.Getter;
import lombok.SneakyThrows;

import javax.sql.DataSource;
import java.sql.BatchUpdateException;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;

import static com.eventsourcing.postgresql.PostgreSQLSerialization.getParameter;
import static com.eventsourcing.postgresql.PostgreSQLSerialization.setValue;

public abstract class PostgreSQLAttributeIndex extends AbstractAttributeIndex {
    protected KeyObjectStore> keyObjectStore;

    protected static  Attribute serializableComparable(Attribute attribute) {
        if (SerializableComparable.class.isAssignableFrom(attribute.getAttributeType())) {
            Class type = SerializableComparable.getType(attribute.getAttributeType());
            @SuppressWarnings("unchecked")
            MultiValueAttribute newAttribute = new SerializableComparableAttribute<>(attribute, type);
            return newAttribute;
        } else {
            return attribute;
        }
    }

    /**
     * Protected constructor, called by subclasses.
     *
     * @param attribute        The attribute on which the index will be built
     * @param supportedQueries The set of {@link Query} types which the subclass implementation supports
     */
    protected PostgreSQLAttributeIndex(Attribute attribute,
                                       Set> supportedQueries) {
        super(attribute, supportedQueries);
    }

    protected abstract DataSource getDataSource();
    protected abstract Layout getLayout();
    protected abstract String getTableName();
    protected abstract TypeHandler getAttributeTypeHandler();
    protected abstract boolean isUnique();

    @SneakyThrows
    public CloseableIterable getDistinctKeys(QueryOptions queryOptions) {
        Connection connection = getDataSource().getConnection();
        PreparedStatement s = connection.prepareStatement("SELECT DISTINCT key FROM " + getTableName() + " ORDER BY key");
        return () -> new PostgreSQLStatementIterator(s, connection, true) {
            @Override public A fetchNext() {
                return (A) PostgreSQLSerialization.getValue(resultSet, new AtomicInteger(1), getAttributeTypeHandler());
            }
        };
    }

    @SneakyThrows
    public Integer getCountForKey(A key, QueryOptions queryOptions) {
        try (Connection connection = getDataSource().getConnection()) {
            try (PreparedStatement s = connection.prepareStatement("SELECT COUNT(key) FROM " + getTableName() + " WHERE " +
                                                                           "key = ?")) {
                setValue(connection, s, 1, getQuantizedValue(key), getAttributeTypeHandler());
                try (java.sql.ResultSet resultSet = s.executeQuery()) {
                    resultSet.next();
                    return resultSet.getInt(1);
                }
            }
        }
    }

    @SneakyThrows
    public Integer getCountOfDistinctKeys(QueryOptions queryOptions) {
        try (Connection connection = getDataSource().getConnection()) {
            try (PreparedStatement s = connection.prepareStatement("SELECT COUNT(DISTINCT key) FROM " + getTableName())) {
                try (java.sql.ResultSet resultSet = s.executeQuery()) {
                    resultSet.next();
                    return resultSet.getInt(1);
                }
            }
        }
    }

    @SneakyThrows
    public CloseableIterable> getStatisticsForDistinctKeys(QueryOptions queryOptions) {
        return getKeyStatisticsForDistinctKeys("ASC");
    }

    @SneakyThrows
    public CloseableIterable> getStatisticsForDistinctKeysDescending(QueryOptions queryOptions) {
        return getKeyStatisticsForDistinctKeys("DESC");
    }


    protected CloseableIterable> getKeyStatisticsForDistinctKeys(String order)
            throws SQLException {Connection connection = getDataSource().getConnection();
        PreparedStatement s = connection.prepareStatement("SELECT DISTINCT key, COUNT(key) FROM " + getTableName() + " " +
                                                                  "GROUP BY key ORDER BY key " + order);
        return new CloseableIterable>() {
            @Override public CloseableIterator> iterator() {
                return new PostgreSQLStatementIterator>(s, connection, true) {
                    @SneakyThrows
                    @Override public KeyStatistics fetchNext() {
                        A key = (A) PostgreSQLSerialization
                                .getValue(resultSet, new AtomicInteger(1), getAttributeTypeHandler());
                        int count = resultSet.getInt(2);
                        return new KeyStatistics<>(key, count);
                    }
                };
            }
        };
    }

    @SneakyThrows
    public CloseableIterable>> getKeysAndValues(QueryOptions queryOptions) {
        return queryKeysAndValues("ASC");
    }

    protected CloseableIterable>> queryKeysAndValues(String order)
            throws SQLException {Connection connection = getDataSource().getConnection();
        PreparedStatement s = connection
                .prepareStatement("SELECT key, value FROM " + getTableName() + " ORDER BY key " + order);
        return new CloseableIterable>>() {
            @Override public CloseableIterator>> iterator() {
                return new PostgreSQLStatementIterator>>(s, connection, true) {
                    @SneakyThrows
                    @Override public KeyValue> fetchNext() {
                        AtomicInteger i = new AtomicInteger(1);
                        A key = (A) PostgreSQLSerialization.getValue(resultSet, i, getAttributeTypeHandler());
                        UUID uuid = UUID.fromString(resultSet.getString(i.get()));
                        return new KeyValueMaterialized<>(key, keyObjectStore.get(uuid));
                    }
                };
            }
        };
    }

    @SneakyThrows
    public CloseableIterable>> getKeysAndValuesDescending(QueryOptions queryOptions) {
        return queryKeysAndValues("DESC");
    }


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

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

    @Override public Index> getEffectiveIndex() {
        return this;
    }

    @Override public boolean addAll(ObjectSet> objectSet, QueryOptions queryOptions) {
        try (CloseableIterator> iterator = objectSet.iterator()) {
            return addAll(iterator, queryOptions);
        }
    }

    @SneakyThrows
    public boolean addAll(Iterator> iterator, QueryOptions queryOptions) {
        try(Connection connection = getDataSource().getConnection()) {
            connection.setAutoCommit(false);
            String insert = "INSERT INTO " + getTableName() + " VALUES (" + getParameter(connection, getAttributeTypeHandler(),
                                                                                    null) + ", ?::UUID) " +
                    (queryOptions.get(OnConflictDo.class) == null ? "" :
                            "ON CONFLICT DO " + queryOptions.get(OnConflictDo.class));
            try (PreparedStatement s = connection.prepareStatement(insert)) {
                while (iterator.hasNext()) {
                    EntityHandle object = iterator.next();
                    Iterator attrIterator = attribute.getValues(object, queryOptions).iterator();
                    while (attrIterator.hasNext()) {
                        int i = 1;
                        A attr = attrIterator.next();
                        i = setValue(connection, s, i, getQuantizedValue(attr), getAttributeTypeHandler());
                        s.setString(i, object.uuid().toString());
                        s.addBatch();
                    }
                }
                try {
                    s.executeBatch();
                } catch (BatchUpdateException e) {
                    connection.rollback();
                    Throwable nextException = e.getCause();
                    if (nextException instanceof PGSQLIntegrityConstraintViolationException) {
                        if (nextException.getMessage().contains("duplicate key value violates unique constraint")) {
                            throw new UniqueIndex.UniqueConstraintViolatedException(nextException.getMessage());
                        } else {
                            throw e;
                        }
                    } else {
                        throw e;
                    }
                }
            }
            connection.commit();
        }

        return true;
    }

    protected void addAll(ObjectStore> objectStore, QueryOptions queryOptions) {
        addAll(objectStore.iterator(queryOptions), queryOptions);
    }

    @SneakyThrows
    @Override public boolean removeAll(ObjectSet> objects, QueryOptions queryOptions) {
        try(Connection connection = getDataSource().getConnection()) {
            String insert = "DELETE FROM " + getTableName() + " WHERE object = ?::UUID";
            try (PreparedStatement s = connection.prepareStatement(insert)) {
                try (CloseableIterator> iterator = objects.iterator()) {
                    while (iterator.hasNext()) {
                        EntityHandle object = iterator.next();
                        s.setString(1, object.uuid().toString());
                        s.addBatch();
                    }
                }
                s.executeBatch();
            }
        }

        return true;
    }

    @SneakyThrows
    @Override public void clear(QueryOptions queryOptions) {
        try(Connection connection = getDataSource().getConnection()) {
            try (PreparedStatement s = connection.prepareStatement("DELETE FROM " + getTableName())) {
                s.executeUpdate();
            }
        }

    }

    @Override public void init(ObjectStore> objectStore, QueryOptions queryOptions) {
        if (objectStore instanceof KeyObjectStore) {
            this.keyObjectStore = (KeyObjectStore>) objectStore;
        } else {
            this.keyObjectStore = new SetKeyObjectStore(objectStore, queryOptions);
        }
        queryOptions.put(OnConflictDo.class, OnConflictDo.NOTHING);
        addAll(objectStore, queryOptions);
    }

    @SneakyThrows
    public CloseableIterable getDistinctKeys(A lowerBound, boolean lowerInclusive, A upperBound,
                                                boolean upperInclusive, QueryOptions queryOptions) {
        return queryDistinctKeys(lowerBound, lowerInclusive, upperBound, upperInclusive, "ASC");
    }

    protected CloseableIterable queryDistinctKeys(A lowerBound, boolean lowerInclusive, A upperBound,
                                                     boolean upperInclusive, String order)
            throws SQLException {Connection connection = getDataSource().getConnection();
        String lowerOp = lowerInclusive ? ">=" : ">";
        String upperOp = upperInclusive ? "<=" : "<";
        String query = "SELECT DISTINCT key FROM " + getTableName() + " WHERE " +
                "key " + lowerOp + " ? AND " +
                "key " + upperOp + " ? " +
                "ORDER BY key " + order;
        PreparedStatement s = connection.prepareStatement(query);
        int i = setValue(connection, s, 1, lowerBound, getAttributeTypeHandler());
        setValue(connection, s, i, upperBound, getAttributeTypeHandler());
        return () -> new PostgreSQLStatementIterator(s, connection, true) {
            @Override public A fetchNext() {
                return (A) PostgreSQLSerialization.getValue(resultSet, new AtomicInteger(1), getAttributeTypeHandler());
            }
        };
    }

    @SneakyThrows
    public CloseableIterable getDistinctKeysDescending(QueryOptions queryOptions) {
        Connection connection = getDataSource().getConnection();
        PreparedStatement s = connection.prepareStatement("SELECT DISTINCT key FROM " + getTableName() + " ORDER BY " +
                                                                  "key DESC");
        return () -> new PostgreSQLStatementIterator(s, connection, true) {
            @Override public A fetchNext() {
                return (A) PostgreSQLSerialization.getValue(resultSet, new AtomicInteger(1), getAttributeTypeHandler());
            }
        };
    }

    @SneakyThrows
    public CloseableIterable getDistinctKeysDescending(A lowerBound, boolean lowerInclusive, A upperBound,
                                                          boolean upperInclusive, QueryOptions queryOptions) {
        return queryDistinctKeys(lowerBound, lowerInclusive, upperBound, upperInclusive, "DESC");
    }


    @SneakyThrows
    public CloseableIterable>> getKeysAndValues(A lowerBound, boolean lowerInclusive,
                                                                            A upperBound, boolean upperInclusive,
                                                                            QueryOptions queryOptions) {
        return queryKeysAndValues(lowerBound, lowerInclusive, upperBound, upperInclusive, queryOptions, "ASC");
    }


    @SneakyThrows
    public CloseableIterable>> getKeysAndValuesDescending(A lowerBound,
                                                                                      boolean lowerInclusive,
                                                                                      A upperBound,
                                                                                      boolean upperInclusive,
                                                                                      QueryOptions queryOptions) {
        return queryKeysAndValues(lowerBound, lowerInclusive, upperBound, upperInclusive, queryOptions, "DESC");
    }

    protected CloseableIterable>> queryKeysAndValues(A lowerBound, boolean lowerInclusive,
                                                                                 A upperBound, boolean upperInclusive,
                                                                                 QueryOptions queryOptionsString,
                                                                                 String order)
            throws SQLException {Connection connection = getDataSource().getConnection();
        String lowerOp = lowerInclusive ? ">=" : ">";
        String upperOp = upperInclusive ? "<=" : "<";
        String sql = "SELECT key, value FROM " + getTableName() +
                " WHERE " +
                "key " + lowerOp + " ? AND " +
                "key " + upperOp + " ? " +
                " ORDER BY key " + order;
        PreparedStatement s = connection.prepareStatement(sql);
        return new CloseableIterable>>() {
            @Override public CloseableIterator>> iterator() {
                return new PostgreSQLStatementIterator>>(s, connection, true) {
                    @SneakyThrows
                    @Override public KeyValue> fetchNext() {
                        AtomicInteger i = new AtomicInteger(1);
                        A key = (A) PostgreSQLSerialization.getValue(resultSet, i, getAttributeTypeHandler());
                        UUID uuid = UUID.fromString(resultSet.getString(i.get()));
                        return new KeyValueMaterialized<>(key, keyObjectStore.get(uuid));
                    }
                };
            }
        };
    }

    @SneakyThrows
    @Override public ResultSet> retrieve(Query> query, QueryOptions queryOptions) {
        Class queryClass = query.getClass();
        if (queryClass.equals(Equal.class)) {
            final Equal, A> equal = (Equal, A>) query;
            Connection connection = getDataSource().getConnection();

            int size;
            A value = ((Equal, A>) query).getValue();

            try(PreparedStatement counter = connection
                    .prepareStatement("SELECT count(object) FROM " + getTableName() + " WHERE key = " + getParameter
                            (connection, getAttributeTypeHandler(), null))) {
                setValue(connection, counter, 1, value,
                         getAttributeTypeHandler());
                try (java.sql.ResultSet resultSet = counter.executeQuery()) {
                    resultSet.next();
                    size = resultSet.getInt(1);
                }
            }

            PreparedStatement s = connection
                    .prepareStatement("SELECT object FROM " + getTableName() + " WHERE key = " +
                                              getParameter(connection, getAttributeTypeHandler(), null));
            setValue(connection, s, 1, value, getAttributeTypeHandler());

            PostgreSQLStatementIterator> iterator = new PostgreSQLStatementIterator>
                    (s, connection, isMutable()) {
                @SneakyThrows
                @Override public EntityHandle fetchNext() {
                    UUID uuid = UUID.fromString(resultSet.getString(1));
                    return keyObjectStore.get(uuid);
                }
            };


            int finalSize = size;
            ResultSet> rs = new MatchingResultSet<>(iterator, equal, queryOptions, finalSize);
            return new CloseableResultSet<>(rs, query, queryOptions);
        } else if (queryClass.equals(Has.class)) {
            final Has, A> has = (Has, A>) query;

            Connection connection = getDataSource().getConnection();

            int size;
            try (PreparedStatement counter = connection
                    .prepareStatement("SELECT count(object) FROM " + getTableName())) {
                try (java.sql.ResultSet resultSet = counter.executeQuery()) {
                    resultSet.next();
                    size = resultSet.getInt(1);
                }
            }

            PreparedStatement s = connection
                    .prepareStatement("SELECT object FROM " + getTableName());

            PostgreSQLStatementIterator> iterator = new PostgreSQLStatementIterator>(s,
                                                                                                                     connection,
                                                                                                                     isMutable()) {
                @SneakyThrows
                @Override public EntityHandle fetchNext() {
                    UUID uuid = UUID.fromString(resultSet.getString(1));
                    return keyObjectStore.get(uuid);
                }
            };

            int finalSize = size;
            ResultSet> rs = new HasResultSet<>(iterator, has, queryOptions, finalSize);
            return new CloseableResultSet<>(rs, query, queryOptions);
        } else {
            throw new IllegalArgumentException("Unsupported query: " + query);
        }
    }

    private static class SerializableComparableAttribute extends MultiValueAttribute {

        private final Attribute attribute;

        public SerializableComparableAttribute(Attribute attribute, Class type) {
            super(attribute.getEffectiveObjectType(), attribute.getObjectType(), (Class) type,
                  attribute.getAttributeName());
            this.attribute = attribute;
        }

        @SuppressWarnings("unchecked")
        @Override public Iterable getValues(Entity object, QueryOptions queryOptions) {
            Iterable iterable = attribute.getValues(new ResolvedEntityHandle(object), queryOptions);
            ArrayList values = new ArrayList<>();
            for (A value : iterable) {
                SerializableComparable value1 = (SerializableComparable) value;
                values.add(value1.getSerializableComparable());
            }
            return values;
        }
    }

    class SetKeyObjectStore implements KeyObjectStore> {

        private final ObjectStore> objectStore;
        private final QueryOptions queryOptions;

        public SetKeyObjectStore(ObjectStore> objectStore, QueryOptions queryOptions) {
            this.objectStore = objectStore;
            this.queryOptions = queryOptions;
        }

        @Override public EntityHandle get(UUID key) {
            CloseableIterator> iterator = objectStore.iterator(queryOptions);
            while (iterator.hasNext()) {
                EntityHandle next = iterator.next();
                if (next.uuid().equals(key)) {
                    return next;
                }
            }
            return null;
        }
    }

    protected class MatchingResultSet>> extends
            ResultSet> {
        private final PostgreSQLStatementIterator> iterator;
        @Getter
        private final T query;
        @Getter
        private final QueryOptions queryOptions;
        private final int finalSize;

        public MatchingResultSet(PostgreSQLStatementIterator> iterator, T query,
                                 QueryOptions queryOptions, int finalSize) {
            this.iterator = iterator;
            this.query = query;
            this.queryOptions = queryOptions;
            this.finalSize = finalSize;
        }

        @Override
        public Iterator> iterator() {
            return iterator;
        }

        @Override
        @SneakyThrows
        public boolean contains(EntityHandle object) {
            try (Connection c = getDataSource().getConnection()) {
                String sql = "SELECT count(key) FROM " + getTableName() + " WHERE object = ?::UUID";
                try (PreparedStatement s = c.prepareStatement(sql)) {
                    try (java.sql.ResultSet resultSet = s.executeQuery()) {
                        resultSet.next();
                        return resultSet.getInt(1) > 0;
                    }
                }
            }
        }

        @Override
        public boolean matches(EntityHandle object) {
            return query.matches(object, queryOptions);
        }

        @Override
        public int getRetrievalCost() {
            return indexRetrievalCost();
        }

        @Override
        public int getMergeCost() {
            return finalSize;
        }

        @Override
        public int size() {
            return finalSize;
        }

        @Override
        public void close() {
            iterator.close();
        }
    }

    private class HasResultSet extends ResultSet> {
        private final PostgreSQLStatementIterator> iterator;
        private final Has, A> has;
        private final QueryOptions queryOptions;
        private final int finalSize;

        public HasResultSet(PostgreSQLStatementIterator> iterator, Has, A> has,
                            QueryOptions queryOptions, int finalSize) {
            this.iterator = iterator;
            this.has = has;
            this.queryOptions = queryOptions;
            this.finalSize = finalSize;
        }

        @Override
        public Iterator> iterator() {
            return iterator;
        }

        @Override
        @SneakyThrows
        public boolean contains(EntityHandle object) {
            try (Connection c = getDataSource().getConnection()) {
                String sql = "SELECT count(key) FROM " + getTableName();
                try (PreparedStatement s = c.prepareStatement(sql)) {
                    try (java.sql.ResultSet resultSet = s.executeQuery()) {
                        resultSet.next();
                        return resultSet.getInt(1) > 0;
                    }
                }
            }
        }

        @Override
        public boolean matches(EntityHandle object) {
            return has.matches(object, queryOptions);
        }

        @Override
        public Query> getQuery() {
            return has;
        }

        @Override
        public QueryOptions getQueryOptions() {
            return queryOptions;
        }

        @Override
        public int getRetrievalCost() {
            return indexRetrievalCost();
        }

        @Override
        public int getMergeCost() {
            return finalSize;
        }

        @Override
        public int size() {
            return finalSize;
        }

        @Override
        public void close() {
            iterator.close();
        }
    }

    protected abstract int indexRetrievalCost();

    protected A getQuantizedValue(A attributeValue) {
        return attributeValue;
    }

    protected enum OnConflictDo {
        UPDATE, NOTHING
    }

}
    

    

    

            

    

            



    
    






    
© 2015 - 2024 Weber Informatics LLC | Privacy Policy