Java tutorial
/* * Copyright 2012-2016 the original author or authors. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. * See the License for the specific language governing permissions and * limitations under the License. */ package org.glowroot.agent.fat.storage.util; import java.sql.Connection; import java.sql.DatabaseMetaData; import java.sql.ResultSet; import java.sql.SQLException; import java.sql.Statement; import java.util.Collection; import java.util.List; import java.util.Locale; import java.util.Map; import java.util.Map.Entry; import java.util.Set; import java.util.TreeMap; import javax.annotation.Nullable; import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ArrayListMultimap; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.google.common.collect.ListMultimap; import com.google.common.collect.Lists; import com.google.common.collect.Maps; import com.google.common.collect.Sets; import org.checkerframework.checker.nullness.qual.PolyNull; import org.checkerframework.checker.tainting.qual.Untainted; import org.immutables.value.Value; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import static com.google.common.base.Preconditions.checkNotNull; import static org.glowroot.agent.fat.storage.util.Checkers.castUntainted; public class Schemas { private static final Logger logger = LoggerFactory.getLogger(Schemas.class); private static final Map<ColumnType, String> typeNames = Maps.newHashMap(); static { // these are type mappings for H2 typeNames.put(ColumnType.VARCHAR, "varchar"); typeNames.put(ColumnType.BIGINT, "bigint"); typeNames.put(ColumnType.BOOLEAN, "boolean"); typeNames.put(ColumnType.VARBINARY, "varbinary"); typeNames.put(ColumnType.DOUBLE, "double"); typeNames.put(ColumnType.AUTO_IDENTITY, "bigint identity"); } private Schemas() { } static void syncTable(@Untainted String tableName, List<Column> columns, Connection connection) throws SQLException { if (!tableExists(tableName, connection)) { createTable(tableName, columns, connection); } else if (tableNeedsUpgrade(tableName, columns, connection)) { logger.warn("upgrading table {}, which unfortunately at this point just means" + " dropping and re-create the table (losing existing data)", tableName); execute("drop table " + tableName, connection); createTable(tableName, columns, connection); } } static void syncIndexes(@Untainted String tableName, ImmutableList<Index> indexes, Connection connection) throws SQLException { ImmutableSet<Index> desiredIndexes = ImmutableSet.copyOf(indexes); Set<Index> existingIndexes = getIndexes(tableName, connection); for (Index index : Sets.difference(existingIndexes, desiredIndexes)) { execute("drop index " + index.name(), connection); } for (Index index : Sets.difference(desiredIndexes, existingIndexes)) { createIndex(tableName, index, connection); } // test the logic existingIndexes = getIndexes(tableName, connection); if (!existingIndexes.equals(desiredIndexes)) { logger.error("the logic in syncIndexes() needs fixing"); } } // useful for upgrades static boolean tableExists(String tableName, Connection connection) throws SQLException { logger.debug("tableExists(): tableName={}", tableName); ResultSet resultSet = getMetaDataTables(connection, tableName); ResultSetCloser closer = new ResultSetCloser(resultSet); try { return resultSet.next(); } catch (Throwable t) { throw closer.rethrow(t); } finally { closer.close(); } } // useful for upgrades static boolean columnExists(String tableName, String columnName, Connection connection) throws SQLException { logger.debug("columnExists(): tableName={}, columnName={}", tableName, columnName); ResultSet resultSet = getMetaDataColumns(connection, tableName, columnName); ResultSetCloser closer = new ResultSetCloser(resultSet); try { return resultSet.next(); } catch (Throwable t) { throw closer.rethrow(t); } finally { closer.close(); } } private static void createTable(@Untainted String tableName, List<Column> columns, Connection connection) throws SQLException { StringBuilder sql = new StringBuilder(); sql.append("create table "); sql.append(tableName); sql.append(" ("); for (int i = 0; i < columns.size(); i++) { if (i > 0) { sql.append(", "); } String sqlTypeName = typeNames.get(columns.get(i).type()); checkNotNull(sqlTypeName, "Unexpected sql type: %s", columns.get(i).type()); sql.append(columns.get(i).name()); sql.append(" "); sql.append(sqlTypeName); if (columns.get(i).primaryKey()) { sql.append(" primary key"); } } sql.append(")"); execute(castUntainted(sql.toString()), connection); if (tableNeedsUpgrade(tableName, columns, connection)) { logger.warn("table {} thinks it still needs to be upgraded, even after it was just" + " upgraded", tableName); } } private static boolean tableNeedsUpgrade(String tableName, List<Column> columns, Connection connection) throws SQLException { // can't use Maps.newTreeMap() because of OpenJDK6 type inference bug // see https://code.google.com/p/guava-libraries/issues/detail?id=635 Map<String, Column> columnMap = new TreeMap<String, Column>(String.CASE_INSENSITIVE_ORDER); for (Column column : columns) { columnMap.put(column.name(), column); } ResultSet resultSet = getMetaDataColumns(connection, tableName, null); ResultSetCloser closer = new ResultSetCloser(resultSet); try { return !columnNamesAndTypesMatch(resultSet, columnMap, connection); } catch (Throwable t) { throw closer.rethrow(t); } finally { closer.close(); } } private static boolean columnNamesAndTypesMatch(ResultSet resultSet, Map<String, Column> columnMap, Connection connection) throws SQLException { while (resultSet.next()) { Column column = columnMap.remove(resultSet.getString("COLUMN_NAME")); if (column == null) { return false; } String typeName = typeNames.get(column.type()); if (typeName == null) { return false; } // this is just to deal with "bigint identity" int index = typeName.indexOf(' '); if (index != -1) { typeName = typeName.substring(0, index); } typeName = convert(connection.getMetaData(), typeName); if (!typeName.equals(resultSet.getString("TYPE_NAME"))) { return false; } } return columnMap.isEmpty(); } @VisibleForTesting static ImmutableSet<Index> getIndexes(String tableName, Connection connection) throws SQLException { ListMultimap</*@Untainted*/ String, /*@Untainted*/ String> indexColumns = ArrayListMultimap.create(); ResultSet resultSet = getMetaDataIndexInfo(connection, tableName); ResultSetCloser closer = new ResultSetCloser(resultSet); try { while (resultSet.next()) { String indexName = checkNotNull(resultSet.getString("INDEX_NAME")); String columnName = checkNotNull(resultSet.getString("COLUMN_NAME")); // hack-ish to skip over primary key constraints which seem to be always // prefixed in H2 by PRIMARY_KEY_ if (!indexName.startsWith("PRIMARY_KEY_")) { indexColumns.put(castUntainted(indexName), castUntainted(columnName)); } } } catch (Throwable t) { throw closer.rethrow(t); } finally { closer.close(); } ImmutableSet.Builder<Index> indexes = ImmutableSet.builder(); for (Entry</*@Untainted*/ String, Collection</*@Untainted*/ String>> entry : indexColumns.asMap() .entrySet()) { String name = entry.getKey().toLowerCase(Locale.ENGLISH); List<String> columns = Lists.newArrayList(); for (String column : entry.getValue()) { columns.add(column.toLowerCase(Locale.ENGLISH)); } indexes.add(ImmutableIndex.of(name, columns)); } return indexes.build(); } private static void createIndex(String tableName, Index index, Connection connection) throws SQLException { StringBuilder sql = new StringBuilder(); sql.append("create index "); sql.append(index.name()); sql.append(" on "); sql.append(tableName); sql.append(" ("); for (int i = 0; i < index.columns().size(); i++) { if (i > 0) { sql.append(", "); } sql.append(index.columns().get(i)); } sql.append(")"); execute(castUntainted(sql.toString()), connection); } private static void execute(@Untainted String sql, Connection connection) throws SQLException { Statement statement = connection.createStatement(); try { statement.execute(sql); } finally { statement.close(); } } private static ResultSet getMetaDataTables(Connection connection, String tableName) throws SQLException { DatabaseMetaData metaData = connection.getMetaData(); return metaData.getTables(null, null, convert(metaData, tableName), null); } private static ResultSet getMetaDataColumns(Connection connection, String tableName, @Nullable String columnName) throws SQLException { DatabaseMetaData metaData = connection.getMetaData(); return metaData.getColumns(null, null, convert(metaData, tableName), convert(metaData, columnName)); } private static ResultSet getMetaDataIndexInfo(Connection connection, String tableName) throws SQLException { DatabaseMetaData metaData = connection.getMetaData(); return metaData.getIndexInfo(null, null, convert(metaData, tableName), false, false); } private static @PolyNull String convert(DatabaseMetaData metaData, @PolyNull String name) throws SQLException { if (name == null) { return null; } if (metaData.storesUpperCaseIdentifiers()) { return name.toUpperCase(Locale.ENGLISH); } else { return name; } } public static enum ColumnType { VARCHAR, BIGINT, BOOLEAN, DOUBLE, VARBINARY, AUTO_IDENTITY; } @Value.Immutable public abstract static class Column { @Value.Parameter abstract String name(); @Value.Parameter abstract ColumnType type(); @Value.Default boolean primaryKey() { return false; } } @Value.Immutable public abstract static class Index { @Value.Parameter abstract @Untainted String name(); @Value.Parameter abstract ImmutableList</*@Untainted*/ String> columns(); } }