org.icij.extract.mysql.MySQLSet Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of extract-cli Show documentation
There is a newer version: 7.4.0
package org.icij.extract.mysql;

import com.mysql.cj.exceptions.MysqlErrorNumbers;

import javax.sql.DataSource;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.*;

public class MySQLSet extends SQLSet {

	private String table;

	public MySQLSet(final DataSource dataSource, final SQLCodec codec, final String table) {
		super(dataSource, codec);
		this.table = table;
	}

	@Override
	public int size() {
		return source.withStatementUnchecked("SELECT COUNT(*) FROM " + table + ";", q -> {
			try (final ResultSet rs = q.executeQuery()) {
				rs.next();
				return rs.getInt(1);
			}
		});
	}

	@Override
	public boolean isEmpty() {
		return source.withStatementUnchecked("SELECT EXISTS(SELECT * FROM " + table + ");", q -> {
			final ResultSet rs = q.executeQuery();

			rs.next();
			return rs.getBoolean(1);
		});
	}

	@Override
	public boolean contains(final Object o) {
		Objects.requireNonNull(o);

		final Map keys = codec.encodeKey(o);
		final Set keySet = keys.keySet();

		final String sql = "SELECT EXISTS(SELECT * FROM " + table + " WHERE " +
				String.join(" AND ", keySet.stream().map(k -> k + " = ?").toArray(String[]::new)) + ";";

		return source.withStatementUnchecked(sql, q -> {
			int i = 1;
			for (String k: keySet) {
				q.setObject(i++, keys.get(k));
			}

			try (final ResultSet rs = q.executeQuery()) {
				rs.next();
				return rs.getBoolean(1);
			}
		});
	}

	@Override
	public boolean add(final E e) {
		Objects.requireNonNull(e);

		final Map values = codec.encodeValue(e);
		values.putAll(codec.encodeKey(e));

		final Set keys = values.keySet();

		final String sql = "INSERT INTO " + table + " (" +
				String.join(", ", keys.toArray(new String[keys.size()])) +
				") VALUES(" +
				String.join(", ", Collections.nCopies(keys.size(), "?")) +
				");";

		return source.withStatementUnchecked(sql, q -> {
			int i = 1;

			for (String key: keys) {
				q.setObject(i++, values.get(key));
			}

			try {
				return q.executeUpdate() > 0;
			} catch (final SQLException exc) {
				if (exc.getErrorCode() == MysqlErrorNumbers.ER_DUP_ENTRY) {
					return false;
				}

				throw exc;
			}
		});
	}

	@Override
	public boolean remove(final Object o) {
		Objects.requireNonNull(o);

		final Map values = codec.encodeKey(o);
		final Set keys = values.keySet();

		final String sql = "DELETE FROM " + table + " WHERE " +
				String.join(" AND ", keys.stream().map(k -> k + " = ?").toArray(String[]::new)) +
				";";

		return source.withStatementUnchecked(sql, q -> {
			int i = 1;

			for (String k: keys) {
				q.setObject(i++, values.get(k));
			}

			return q.executeUpdate() > 0;
		});
	}

	@Override
	public void clear() {
		source.withStatementUnchecked("DELETE FROM " + table + ";",
				(CheckedFunction) PreparedStatement::executeUpdate);
	}

	@Override
	public Iterator iterator() {
		return new MySQLIterator<>(source, codec, table);
	}
}