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

org.vertexium.sql.collections.SqlMap Maven / Gradle / Ivy

The newest version!
package org.vertexium.sql.collections;

import org.skife.jdbi.v2.*;
import org.skife.jdbi.v2.tweak.ResultSetMapper;
import org.skife.jdbi.v2.util.ByteArrayMapper;
import org.skife.jdbi.v2.util.IntegerMapper;
import org.skife.jdbi.v2.util.StringMapper;
import org.vertexium.VertexiumSerializer;
import org.vertexium.util.CloseableUtils;
import org.vertexium.util.VertexiumLogger;
import org.vertexium.util.VertexiumLoggerFactory;

import javax.sql.DataSource;
import java.io.Closeable;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.*;

@SuppressWarnings("NullableProblems")
public class SqlMap extends AbstractMap {
    private static final VertexiumLogger LOGGER = VertexiumLoggerFactory.getLogger(SqlMap.class);

    protected final String tableName;
    protected final String keyColumnName;
    protected final String valueColumnName;
    private final DBI dbi;
    private final VertexiumSerializer serializer;
    private final ResultSetMapper> entrySetMapper;
    private Object storableContext;

    public SqlMap(String tableName, String keyColumnName, String valueColumnName, DataSource dataSource,
                  VertexiumSerializer serializer) {
        this.tableName = tableName;
        this.keyColumnName = keyColumnName;
        this.valueColumnName = valueColumnName;
        this.dbi = new DBI(dataSource);
        this.serializer = serializer;
        this.entrySetMapper = new ResultSetMapper>() {
            public MapEntry map(int index, ResultSet rs, StatementContext ctx) throws SQLException {
                String key = rs.getString(SqlMap.this.keyColumnName);
                byte[] value = rs.getBytes(SqlMap.this.valueColumnName);
                return new MapEntry<>(key, value);
            }
        };
    }

    public void setStorableContext(Object storableContext) {
        this.storableContext = storableContext;
    }

    @Override
    public Set> entrySet() {
        final Handle handle = dbi.open();
        final Query> query = handle
                .createQuery(String.format(
                        "select %s, %s from %s order by %s", keyColumnName, valueColumnName, tableName, keyColumnName))
                .map(entrySetMapper);

        return new IteratingSet>() {
            @Override
            public Iterator> createIterator() {
                return new QueryResultIterator, MapEntry>(query, handle) {
                    @Override
                    public Entry next() {
                        MapEntry stringifiedEntry = resultIterator.next();
                        String key = stringifiedEntry.getKey();
                        T value = withContainer(serializer.bytesToObject(stringifiedEntry.getValue()));
                        return new MapEntry<>(key, value);
                    }
                };
            }
        };
    }

    @Override
    public Set keySet() {
        final Handle handle = dbi.open();
        final Query query = handle
                .createQuery(String.format("select %s from %s order by %s", keyColumnName, tableName, keyColumnName))
                .map(StringMapper.FIRST);

        return new IteratingSet() {
            @Override
            public Iterator createIterator() {
                return new QueryResultIterator<>(query, handle);
            }
        };
    }

    @Override
    public Collection values() {
        final Handle handle = dbi.open();
        final Query query = handle
                .createQuery(String.format("select %s from %s order by %s", valueColumnName, tableName, keyColumnName))
                .map(ByteArrayMapper.FIRST);

        return new IteratingSet() {
            @Override
            public Iterator createIterator() {
                return new QueryResultIterator(query, handle) {
                    @Override
                    public T next() {
                        return withContainer(serializer.bytesToObject(resultIterator.next()));
                    }
                };
            }
        };
    }

    @Override
    public boolean containsKey(Object key) {
        try (Handle handle = dbi.open()){
            return handle
                    .createQuery(String.format(
                            "select count(*) %s from %s where %s = ?", keyColumnName, tableName, keyColumnName))
                    .bind(0, key)
                    .map(IntegerMapper.FIRST)
                    .first() > 0;
        }
    }

    @Override
    public boolean containsValue(Object value) {
        try (Handle handle = dbi.open()) {
            return handle
                    .createQuery(String.format(
                            "select count(*) %s from %s where %s = ?", valueColumnName, tableName, valueColumnName))
                    .bind(0, serializer.objectToBytes(value))
                    .map(IntegerMapper.FIRST)
                    .first() > 0;
        }
    }

    @Override
    public int size() {
        try (Handle handle = dbi.open()) {
            return handle
                    .createQuery(String.format("select count(*) from %s", tableName))
                    .map(IntegerMapper.FIRST)
                    .first();
        }
    }

    @Override
    public void clear() {
        try (Handle handle = dbi.open()) {
            handle.execute(String.format("delete from %s", tableName));
        }
    }

    @Override
    public T remove(Object key) {
        T value = get(key);
        try (Handle handle = dbi.open()) {
            handle.execute(String.format("delete from %s where %s = ?", tableName, keyColumnName), key);
            return withoutContainer(value);
        }
    }

    @Override
    public T get(Object key) {
        try (Handle handle = dbi.open()) {
            return withContainer(serializer.bytesToObject(handle
                    .createQuery(String.format(
                            "select %s from %s where %s = ?", valueColumnName, tableName, keyColumnName))
                    .bind(0, key)
                    .map(ByteArrayMapper.FIRST)
                    .first()));
        }
    }

    @Override
    public T put(String key, T value) {
        byte[] byteArrayValue = serializer.objectToBytes(withContainer(value));
        T previous = get(key);
        try (Handle handle = dbi.open()) {
            if (previous == null) {
                handle.execute(String.format(
                        "insert into %s (%s, %s) values (?, ?)", tableName, keyColumnName, valueColumnName),
                        key, byteArrayValue);
            } else {
                handle.execute(String.format(
                        "update %s set %s = ? where %s = ?", tableName, valueColumnName, keyColumnName),
                        byteArrayValue, key);
            }
            updateAdditionalColumns(handle, key, value);
        }

        return withoutContainer(previous);
    }

    private void updateAdditionalColumns(Handle handle, String key, T value) {
        Map additional = additionalColumns(key, value);
        if (!additional.isEmpty()) {
            StringBuilder updateSql = new StringBuilder(String.format("update %s set ", tableName));
            List positionalParams = new ArrayList<>();
            boolean first = true;
            for (Entry column : additional.entrySet()) {
                if (first) {
                    updateSql.append(String.format("%s = ?", column.getKey()));
                    first = false;
                } else {
                    updateSql.append(String.format(", %s = ?", column.getKey()));
                }
                positionalParams.add(column.getValue());
            }
            updateSql.append(String.format(" where %s = ?", keyColumnName));
            positionalParams.add(key);
            handle.execute(updateSql.toString(), positionalParams.toArray());
        }
    }

    public Iterator query(String where, Object... positionalParams) {
        final Handle handle = dbi.open();
        Query> query1 = handle.createQuery(String.format(
                "select %s from %s where %s order by %s", valueColumnName, tableName, where, keyColumnName));
        int i = 0;
        for (Object param : positionalParams) {
            query1 = query1.bind(i++, param);
        }
        final Query query2 = query1.map(ByteArrayMapper.FIRST);

        return new QueryResultIterator(query2, handle) {
            @Override
            public T next() {
                return withContainer(serializer.bytesToObject(resultIterator.next()));
            }
        };
    }

    public Iterator query(String where, Map namedParams) {
        final Handle handle = dbi.open();
        Query> query1 = handle.createQuery(String.format(
                "select %s from %s where %s order by %s", valueColumnName, tableName, where, keyColumnName));
        for (Map.Entry param : namedParams.entrySet()) {
            query1 = query1.bind(param.getKey(), param.getValue());
        }
        final Query query2 = query1.map(ByteArrayMapper.FIRST);

        return new QueryResultIterator(query2, handle) {
            @Override
            public T next() {
                return withContainer(serializer.bytesToObject(resultIterator.next()));
            }
        };
    }

    @SuppressWarnings("unused")
    protected Map additionalColumns(String key, T value) {
        // subclasses can override to supply additional column data to be stored, for supporting custom queries.
        return Collections.emptyMap();
    }

    @SuppressWarnings("unchecked")
    private T withContainer(T value) {
        if (value instanceof Storable) {
            ((Storable) value).setContainer(this, storableContext);
        }
        return value;
    }

    private T withoutContainer(T value) {
        if (value instanceof Storable) {
            ((Storable) value).setContainer(null, null);
        }
        return value;
    }

    private abstract class IteratingSet extends AbstractSet implements Closeable {
        private Iterator iterator;

        @Override
        public final Iterator iterator() {
            if (iterator != null) throw new IllegalStateException("can't ask for iterator more than once");
            iterator = createIterator();
            return iterator;
        }

        @Override
        public int size() {
            return SqlMap.this.size();
        }

        @Override
        public void clear() {
            SqlMap.this.clear();
        }

        @Override
        public boolean contains(Object v) {
            return SqlMap.this.containsValue(v);
        }

        protected abstract Iterator createIterator();

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

    private class QueryResultIterator implements Iterator, Closeable {
        protected final Handle handle;
        protected final ResultIterator resultIterator;
        private final Throwable creatorTrace;
        private boolean closed = false;

        QueryResultIterator(Query query, Handle handle) {
            this.creatorTrace = new Throwable();
            this.handle = handle;
            this.resultIterator = query.iterator();
        }

        @Override
        public boolean hasNext() {
            boolean hasNext = resultIterator.hasNext();
            if (!hasNext) {
                close();
            }
            return hasNext;
        }

        @SuppressWarnings("unchecked")
        @Override
        public E next() {
            return (E) resultIterator.next();
        }

        @Override
        public void remove() {
            resultIterator.remove();
        }

        @Override
        protected void finalize() throws Throwable {
            super.finalize();
            // If the iterator isn't completely consumed, then this provides a fallback to close the
            // handle, which holds an open JDBC connection.
            if (!closed) {
                LOGGER.warn("closing QueryResultIterator handle from finalizer", creatorTrace);
                close();
            }
        }

        public void close() {
            handle.close();
            closed = true;
        }
    }
}