opsu-dance/src/itdelatrisu/opsu/db/BeatmapDB.java

/*
 * opsu! - an open-source osu! client
 * Copyright (C) 2014, 2015 Jeffrey Han
 *
 * opsu! is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * opsu! is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with opsu!.  If not, see <http://www.gnu.org/licenses/>.
 */

package itdelatrisu.opsu.db;

import itdelatrisu.opsu.ErrorHandler;
import itdelatrisu.opsu.Options;
import itdelatrisu.opsu.beatmap.Beatmap;
import itdelatrisu.opsu.beatmap.BeatmapParser;

import java.io.File;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import org.newdawn.slick.util.Log;

/**
 * Handles connections and queries with the cached beatmap database.
 */
public class BeatmapDB {
	/**
	 * Current database version.
	 * This value should be changed whenever the database format changes.
	 */
	private static final String DATABASE_VERSION = "2015-06-11";

	/** Minimum batch size ratio ({@code batchSize/cacheSize}) to invoke batch loading. */
	private static final float LOAD_BATCH_MIN_RATIO = 0.2f;

	/** Minimum batch size to invoke batch insertion. */
	private static final int INSERT_BATCH_MIN = 100;

	/** Beatmap loading flags. */
	public static final int LOAD_NONARRAY = 1, LOAD_ARRAY = 2, LOAD_ALL = 3;

	/** Database connection. */
	private static Connection connection;

	/** Query statements. */
	private static PreparedStatement insertStmt, selectStmt, deleteMapStmt, deleteGroupStmt, updateSizeStmt;

	/** Current size of beatmap cache table. */
	private static int cacheSize = -1;

	// This class should not be instantiated.
	private BeatmapDB() {}

	/**
	 * Initializes the database connection.
	 */
	public static void init() {
		// create a database connection
		connection = DBController.createConnection(Options.BEATMAP_DB.getPath());
		if (connection == null)
			return;

		// create the database
		createDatabase();

		// prepare sql statements (used below)
		try {
			updateSizeStmt = connection.prepareStatement("REPLACE INTO info (key, value) VALUES ('size', ?)");
		} catch (SQLException e) {
			ErrorHandler.error("Failed to prepare beatmap statements.", e, true);
		}

		// retrieve the cache size
		getCacheSize();

		// check the database version
		checkVersion();

		// prepare sql statements (not used here)
		try {
			insertStmt = connection.prepareStatement(
				"INSERT INTO beatmaps VALUES (" +
				"?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, " +
				"?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"
			);
			selectStmt = connection.prepareStatement("SELECT * FROM beatmaps WHERE dir = ? AND file = ?");
			deleteMapStmt = connection.prepareStatement("DELETE FROM beatmaps WHERE dir = ? AND file = ?");
			deleteGroupStmt = connection.prepareStatement("DELETE FROM beatmaps WHERE dir = ?");
		} catch (SQLException e) {
			ErrorHandler.error("Failed to prepare beatmap statements.", e, true);
		}
	}

	/**
	 * Creates the database, if it does not exist.
	 */
	private static void createDatabase() {
		try (Statement stmt = connection.createStatement()) {
			String sql =
				"CREATE TABLE IF NOT EXISTS beatmaps (" +
					"dir TEXT, file TEXT, lastModified INTEGER, " +
					"MID INTEGER, MSID INTEGER, " +
					"title TEXT, titleUnicode TEXT, artist TEXT, artistUnicode TEXT, " +
					"creator TEXT, version TEXT, source TEXT, tags TEXT, " +
					"circles INTEGER, sliders INTEGER, spinners INTEGER, " +
					"hp REAL, cs REAL, od REAL, ar REAL, sliderMultiplier REAL, sliderTickRate REAL, " +
					"bpmMin INTEGER, bpmMax INTEGER, endTime INTEGER, " +
					"audioFile TEXT, audioLeadIn INTEGER, previewTime INTEGER, countdown INTEGER, sampleSet TEXT, stackLeniency REAL, " +
					"mode INTEGER, letterboxInBreaks BOOLEAN, widescreenStoryboard BOOLEAN, epilepsyWarning BOOLEAN, " +
					"bg TEXT, sliderBorder TEXT, timingPoints TEXT, breaks TEXT, combo TEXT, " +
					"md5hash TEXT" +
				"); " +
				"CREATE TABLE IF NOT EXISTS info (" +
					"key TEXT NOT NULL UNIQUE, value TEXT" +
				"); " +
				"CREATE INDEX IF NOT EXISTS idx ON beatmaps (dir, file); " +

				// extra optimizations
				"PRAGMA locking_mode = EXCLUSIVE; " +
				"PRAGMA journal_mode = WAL;";
			stmt.executeUpdate(sql);

			// set the version key, if empty
			sql = String.format("INSERT OR IGNORE INTO info(key, value) VALUES('version', '%s')", DATABASE_VERSION);
			stmt.executeUpdate(sql);
		} catch (SQLException e) {
			ErrorHandler.error("Could not create beatmap database.", e, true);
		}
	}

	/**
	 * Checks the stored table version, clears the beatmap database if different
	 * from the current version, then updates the version field.
	 */
	private static void checkVersion() {
		try (Statement stmt = connection.createStatement()) {
			// get the stored version
			String sql = "SELECT value FROM info WHERE key = 'version'";
			ResultSet rs = stmt.executeQuery(sql);
			String version = (rs.next()) ? rs.getString(1) : "";
			rs.close();

			// if different from current version, clear the database
			if (!version.equals(DATABASE_VERSION)) {
				clearDatabase();

				// update version
				PreparedStatement ps = connection.prepareStatement("REPLACE INTO info (key, value) VALUES ('version', ?)");
				ps.setString(1, DATABASE_VERSION);
				ps.executeUpdate();
				ps.close();
			}
		} catch (SQLException e) {
			ErrorHandler.error("Beatmap database version checks failed.", e, true);
		}
	}

	/**
	 * Retrieves the size of the beatmap cache from the 'info' table.
	 */
	private static void getCacheSize() {
		try (Statement stmt = connection.createStatement()) {
			String sql = "SELECT value FROM info WHERE key = 'size'";
			ResultSet rs = stmt.executeQuery(sql);
			try {
				cacheSize = (rs.next()) ? Integer.parseInt(rs.getString(1)) : 0;
			} catch (NumberFormatException e) {
				cacheSize = 0;
			}
			rs.close();
		} catch (SQLException e) {
			ErrorHandler.error("Could not get beatmap cache size.", e, true);
		}
	}

	/**
	 * Updates the size of the beatmap cache in the 'info' table.
	 */
	private static void updateCacheSize() {
		if (connection == null)
			return;

		try {
			updateSizeStmt.setString(1, Integer.toString(Math.max(cacheSize, 0)));
			updateSizeStmt.executeUpdate();
		} catch (SQLException e) {
			ErrorHandler.error("Could not update beatmap cache size.", e, true);
		}
	}

	/**
	 * Clears the database.
	 */
	public static void clearDatabase() {
		if (connection == null)
			return;

		// drop the table, then recreate it
		try (Statement stmt = connection.createStatement()) {
			String sql = "DROP TABLE beatmaps";
			stmt.executeUpdate(sql);
			cacheSize = 0;
			updateCacheSize();
		} catch (SQLException e) {
			ErrorHandler.error("Could not drop beatmap database.", e, true);
		}
		createDatabase();
	}

	/**
	 * Adds the beatmap to the database.
	 * @param beatmap the beatmap
	 */
	public static void insert(Beatmap beatmap) {
		if (connection == null)
			return;

		try {
			setStatementFields(insertStmt, beatmap);
			cacheSize += insertStmt.executeUpdate();
			updateCacheSize();
		} catch (SQLException e) {
			ErrorHandler.error("Failed to add beatmap to database.", e, true);
		}
	}

	/**
	 * Adds the beatmaps to the database in a batch.
	 * @param batch a list of beatmaps
	 */
	public static void insert(List<Beatmap> batch) {
		if (connection == null)
			return;

		try (Statement stmt = connection.createStatement()) {
			// turn off auto-commit mode
			boolean autoCommit = connection.getAutoCommit();
			connection.setAutoCommit(false);

			// drop indexes
			boolean recreateIndexes = (batch.size() >= INSERT_BATCH_MIN);
			if (recreateIndexes) {
				String sql = "DROP INDEX IF EXISTS idx";
				stmt.executeUpdate(sql);
			}

			// batch insert
			for (Beatmap beatmap : batch) {
				try {
					setStatementFields(insertStmt, beatmap);
				} catch (SQLException e) {
					Log.error(String.format("Failed to insert map '%s' into database.", beatmap.getFile().getPath()), e);
					continue;
				}
				insertStmt.addBatch();
			}
			int[] results = insertStmt.executeBatch();
			for (int i = 0; i < results.length; i++) {
				if (results[i] > 0)
					cacheSize += results[i];
			}

			// re-create indexes
			if (recreateIndexes) {
				String sql = "CREATE INDEX idx ON beatmaps (dir, file)";
				stmt.executeUpdate(sql);
			}

			// restore previous auto-commit mode
			connection.commit();
			connection.setAutoCommit(autoCommit);

			// update cache size
			updateCacheSize();
		} catch (SQLException e) {
			ErrorHandler.error("Failed to add beatmaps to database.", e, true);
		}
	}

	/**
	 * Sets all statement fields using a given beatmap.
	 * @param stmt the statement to set fields for
	 * @param beatmap the beatmap
	 * @throws SQLException
	 */
	private static void setStatementFields(PreparedStatement stmt, Beatmap beatmap)
			throws SQLException {
		try {
			stmt.setString(1, beatmap.getFile().getParentFile().getName());
			stmt.setString(2, beatmap.getFile().getName());
			stmt.setLong(3, beatmap.getFile().lastModified());
			stmt.setInt(4, beatmap.beatmapID);
			stmt.setInt(5, beatmap.beatmapSetID);
			stmt.setString(6, beatmap.title);
			stmt.setString(7, beatmap.titleUnicode);
			stmt.setString(8, beatmap.artist);
			stmt.setString(9, beatmap.artistUnicode);
			stmt.setString(10, beatmap.creator);
			stmt.setString(11, beatmap.version);
			stmt.setString(12, beatmap.source);
			stmt.setString(13, beatmap.tags);
			stmt.setInt(14, beatmap.hitObjectCircle);
			stmt.setInt(15, beatmap.hitObjectSlider);
			stmt.setInt(16, beatmap.hitObjectSpinner);
			stmt.setFloat(17, beatmap.HPDrainRate);
			stmt.setFloat(18, beatmap.circleSize);
			stmt.setFloat(19, beatmap.overallDifficulty);
			stmt.setFloat(20, beatmap.approachRate);
			stmt.setFloat(21, beatmap.sliderMultiplier);
			stmt.setFloat(22, beatmap.sliderTickRate);
			stmt.setInt(23, beatmap.bpmMin);
			stmt.setInt(24, beatmap.bpmMax);
			stmt.setInt(25, beatmap.endTime);
			stmt.setString(26, beatmap.audioFilename.getName());
			stmt.setInt(27, beatmap.audioLeadIn);
			stmt.setInt(28, beatmap.previewTime);
			stmt.setByte(29, beatmap.countdown);
			stmt.setString(30, beatmap.sampleSet);
			stmt.setFloat(31, beatmap.stackLeniency);
			stmt.setByte(32, beatmap.mode);
			stmt.setBoolean(33, beatmap.letterboxInBreaks);
			stmt.setBoolean(34, beatmap.widescreenStoryboard);
			stmt.setBoolean(35, beatmap.epilepsyWarning);
			stmt.setString(36, (beatmap.bg == null) ? null : beatmap.bg.getName());
			stmt.setString(37, beatmap.sliderBorderToString());
			stmt.setString(38, beatmap.timingPointsToString());
			stmt.setString(39, beatmap.breaksToString());
			stmt.setString(40, beatmap.comboToString());
			stmt.setString(41, beatmap.md5Hash);
		} catch (SQLException e) {
			throw e;
		} catch (Exception e) {
			throw new SQLException(e);
		}
	}

	/**
	 * Loads beatmap fields from the database.
	 * @param beatmap the beatmap
	 * @param flag whether to load all fields (LOAD_ALL), non-array
	 *        fields (LOAD_NONARRAY), or array fields (LOAD_ARRAY)
	 */
	public static void load(Beatmap beatmap, int flag) {
		if (connection == null)
			return;

		try {
			selectStmt.setString(1, beatmap.getFile().getParentFile().getName());
			selectStmt.setString(2, beatmap.getFile().getName());
			ResultSet rs = selectStmt.executeQuery();
			if (rs.next()) {
				if ((flag & LOAD_NONARRAY) > 0)
					setBeatmapFields(rs, beatmap);
				if ((flag & LOAD_ARRAY) > 0)
					setBeatmapArrayFields(rs, beatmap);
			}
			rs.close();
		} catch (SQLException e) {
			ErrorHandler.error("Failed to load Beatmap from database.", e, true);
		}
	}

	/**
	 * Loads Beatmap fields from the database in a batch.
	 * @param batch a list of beatmaps
	 * @param flag whether to load all fields (LOAD_ALL), non-array
	 *        fields (LOAD_NONARRAY), or array fields (LOAD_ARRAY)
	 */
	public static void load(List<Beatmap> batch, int flag) {
		if (connection == null)
			return;

		// batch size too small
		int size = batch.size();
		if (size < cacheSize * LOAD_BATCH_MIN_RATIO) {
			for (Beatmap beatmap : batch)
				load(beatmap, flag);
			return;
		}

		try (Statement stmt = connection.createStatement()) {
			// create map
			HashMap<String, HashMap<String, Beatmap>> map = new HashMap<String, HashMap<String, Beatmap>>();
			for (Beatmap beatmap : batch) {
				String parent = beatmap.getFile().getParentFile().getName();
				String name = beatmap.getFile().getName();
				HashMap<String, Beatmap> m = map.get(parent);
				if (m == null) {
					m = new HashMap<String, Beatmap>();
					map.put(parent, m);
				}
				m.put(name, beatmap);
			}

			// iterate through database to load beatmaps
			int count = 0;
			stmt.setFetchSize(100);
			String sql = "SELECT * FROM beatmaps";
			ResultSet rs = stmt.executeQuery(sql);
			while (rs.next()) {
				String parent = rs.getString(1);
				HashMap<String, Beatmap> m = map.get(parent);
				if (m != null) {
					String name = rs.getString(2);
					Beatmap beatmap = m.get(name);
					if (beatmap != null) {
						try {
							if ((flag & LOAD_NONARRAY) > 0)
								setBeatmapFields(rs, beatmap);
							if ((flag & LOAD_ARRAY) > 0)
								setBeatmapArrayFields(rs, beatmap);
						} catch (SQLException e) {
							Log.error(String.format("Failed to load map '%s/%s' from database.", parent, name), e);
						}
						if (++count >= size)
							break;
					}
				}
			}
			rs.close();
		} catch (SQLException e) {
			ErrorHandler.error("Failed to load beatmaps from database.", e, true);
		}
	}

	/**
	 * Sets all beatmap non-array fields using a given result set.
	 * @param rs the result set containing the fields
	 * @param beatmap the beatmap
	 * @throws SQLException
	 */
	private static void setBeatmapFields(ResultSet rs, Beatmap beatmap) throws SQLException {
		try {
			File dir = beatmap.getFile().getParentFile();
			beatmap.beatmapID = rs.getInt(4);
			beatmap.beatmapSetID = rs.getInt(5);
			beatmap.title = BeatmapParser.getDBString(rs.getString(6));
			beatmap.titleUnicode = BeatmapParser.getDBString(rs.getString(7));
			beatmap.artist = BeatmapParser.getDBString(rs.getString(8));
			beatmap.artistUnicode = BeatmapParser.getDBString(rs.getString(9));
			beatmap.creator = BeatmapParser.getDBString(rs.getString(10));
			beatmap.version = BeatmapParser.getDBString(rs.getString(11));
			beatmap.source = BeatmapParser.getDBString(rs.getString(12));
			beatmap.tags = BeatmapParser.getDBString(rs.getString(13));
			beatmap.hitObjectCircle = rs.getInt(14);
			beatmap.hitObjectSlider = rs.getInt(15);
			beatmap.hitObjectSpinner = rs.getInt(16);
			beatmap.HPDrainRate = rs.getFloat(17);
			beatmap.circleSize = rs.getFloat(18);
			beatmap.overallDifficulty = rs.getFloat(19);
			beatmap.approachRate = rs.getFloat(20);
			beatmap.sliderMultiplier = rs.getFloat(21);
			beatmap.sliderTickRate = rs.getFloat(22);
			beatmap.bpmMin = rs.getInt(23);
			beatmap.bpmMax = rs.getInt(24);
			beatmap.endTime = rs.getInt(25);
			beatmap.audioFilename = new File(dir, BeatmapParser.getDBString(rs.getString(26)));
			beatmap.audioLeadIn = rs.getInt(27);
			beatmap.previewTime = rs.getInt(28);
			beatmap.countdown = rs.getByte(29);
			beatmap.sampleSet = BeatmapParser.getDBString(rs.getString(30));
			beatmap.stackLeniency = rs.getFloat(31);
			beatmap.mode = rs.getByte(32);
			beatmap.letterboxInBreaks = rs.getBoolean(33);
			beatmap.widescreenStoryboard = rs.getBoolean(34);
			beatmap.epilepsyWarning = rs.getBoolean(35);
			String bg = rs.getString(36);
			if (bg != null)
				beatmap.bg = new File(dir, BeatmapParser.getDBString(bg));
			beatmap.sliderBorderFromString(rs.getString(37));
			beatmap.md5Hash = rs.getString(41);
		} catch (SQLException e) {
			throw e;
		} catch (Exception e) {
			throw new SQLException(e);
		}
	}

	/**
	 * Sets all Beatmap array fields using a given result set.
	 * @param rs the result set containing the fields
	 * @param beatmap the beatmap
	 * @throws SQLException
	 */
	private static void setBeatmapArrayFields(ResultSet rs, Beatmap beatmap) throws SQLException {
		try {
			beatmap.timingPointsFromString(rs.getString(38));
			beatmap.breaksFromString(rs.getString(39));
			beatmap.comboFromString(rs.getString(40));
		} catch (SQLException e) {
			throw e;
		} catch (Exception e) {
			throw new SQLException(e);
		}
	}

	/**
	 * Returns a map of file paths ({dir}/{file}) to last modified times, or
	 * null if any error occurred.
	 */
	public static Map<String, Long> getLastModifiedMap() {
		if (connection == null)
			return null;

		try (Statement stmt = connection.createStatement()) {
			Map<String, Long> map = new HashMap<String, Long>();
			String sql = "SELECT dir, file, lastModified FROM beatmaps";
			ResultSet rs = stmt.executeQuery(sql);
			stmt.setFetchSize(100);
			while (rs.next()) {
				String path = String.format("%s/%s", rs.getString(1), rs.getString(2));
				long lastModified = rs.getLong(3);
				map.put(path, lastModified);
			}
			rs.close();
			return map;
		} catch (SQLException e) {
			ErrorHandler.error("Failed to get last modified map from database.", e, true);
			return null;
		}
	}

	/**
	 * Deletes the beatmap entry from the database.
	 * @param dir the directory
	 * @param file the file
	 */
	public static void delete(String dir, String file) {
		if (connection == null)
			return;

		try {
			deleteMapStmt.setString(1, dir);
			deleteMapStmt.setString(2, file);
			cacheSize -= deleteMapStmt.executeUpdate();
			updateCacheSize();
		} catch (SQLException e) {
			ErrorHandler.error("Failed to delete beatmap entry from database.", e, true);
		}
	}

	/**
	 * Deletes the beatmap group entry from the database.
	 * @param dir the directory
	 */
	public static void delete(String dir) {
		if (connection == null)
			return;

		try {
			deleteGroupStmt.setString(1, dir);
			cacheSize -= deleteGroupStmt.executeUpdate();
			updateCacheSize();
		} catch (SQLException e) {
			ErrorHandler.error("Failed to delete beatmap group entry from database.", e, true);
		}
	}

	/**
	 * Closes the connection to the database.
	 */
	public static void closeConnection() {
		if (connection == null)
			return;

		try {
			insertStmt.close();
			selectStmt.close();
			deleteMapStmt.close();
			deleteGroupStmt.close();
			updateSizeStmt.close();
			connection.close();
			connection = null;
		} catch (SQLException e) {
			ErrorHandler.error("Failed to close beatmap database.", e, true);
		}
	}
}