獲取 MySQL, Oracle, MongoDB 數據庫中所有表的數據量

關系型數據庫

package com.seliote.databasecountsnapshot.tool;

import com.seliote.databasecountsnapshot.pojo.RdbmsConfig;
import com.seliote.databasecountsnapshot.pojo.TableCount;
import lombok.extern.slf4j.Slf4j;

import java.sql.*;
import java.util.Enumeration;
import java.util.LinkedList;
import java.util.List;

/**
 * RDBMS 相關操作工具類
 *
 * @author LiYangDi
 * @since 2019/12/13
 */
@Slf4j
public class RdbmsOperate {

    /**
     * 注冊數據庫驅動
     */
    public static void registerDrivers() throws ClassNotFoundException {
        log.info("Register drivers");
        // Oracle 驅動
        Class.forName("oracle.jdbc.OracleDriver");
        // MySQL 驅動
        //Class.forName("com.mysql.jdbc.Driver");
        Class.forName("com.mysql.cj.jdbc.Driver");
    }

    /**
     * 注銷數據庫驅動
     */
    public static void unregisterDrivers() {
        log.info("Unregister drivers");
        // 獲取加載器，下方需要判斷驅動是否是應用自身加載的
        ClassLoader classLoader = Thread.currentThread().getContextClassLoader();
        // 獲取已加載的所有驅動
        Enumeration<Driver> driverEnumeration = DriverManager.getDrivers();
        while (driverEnumeration.hasMoreElements()) {
            Driver driver = driverEnumeration.nextElement();
            // 需要直接判斷是否是同一個對象，所以使用了 == 而不是 equals(.)
            if (driver.getClass().getClassLoader() == classLoader) {
                try {
                    // 注銷驅動
                    DriverManager.deregisterDriver(driver);
                } catch (SQLException exp) {
                    log.warn("Exception when unregister drivers: " + exp);
                }
            }
        }
    }

    /**
     * 獲取 RDBMS 數據源連接
     *
     * @return RDBMS 數據源
     */
    public static Connection getConnection(RdbmsConfig config) throws SQLException {
        log.info("Attempt to get DataSource list for: " + config);
        return DriverManager.getConnection(config.getUrl(),
                config.getUsername(),
                config.getPassword());
    }

    /**
     * 關閉所有 RDBMS 數據源連接
     *
     * @param connection RDBMS 數據源連接
     */
    public static void closeConnection(Connection connection) throws SQLException {
        log.info("Close connection: " + connection);
        connection.close();
    }

    /**
     * 獲取連接的表中所有數據量
     *
     * @param connection 連接對象
     * @return 表中所有數據量
     */
    public static List<TableCount> getTableCounts(Connection connection, String username) throws SQLException {
        log.info("Getting RDBMS tables counts");
        List<TableCount> tableCounts = new LinkedList<>();
        DatabaseMetaData databaseMetaData = connection.getMetaData();
        // 數據庫表元數據
        try (ResultSet resultSet = databaseMetaData.getTables(connection.getCatalog(),
                // 這個只對 Oracle 有用，設置成大寫的用戶名即可
                username.toUpperCase(),
                null,
                new String[]{"TABLE"})) {
            // 獲取表名
            while (resultSet.next()) {
                // 只獲取表類型
                String tableType = resultSet.getString("TABLE_TYPE");
                if (!tableType.equals("TABLE")) {
                    log.info("Table " + resultSet.getString("TABLE_NAME")
                            + " type is not `TABLE`");
                    continue;
                }
                // Oracle 下庫名是 null
                String databaseName = resultSet.getString("TABLE_CAT");
                String tableName = resultSet.getString("TABLE_NAME");
                long count = -1L;
                // 獲取行數
                //noinspection SqlNoDataSourceInspection
                try (PreparedStatement statement
                             = connection.prepareStatement("SELECT COUNT(*) FROM " + tableName)) {
                    try (ResultSet countResultSet = statement.executeQuery()) {
                        while (countResultSet.next()) {
                            count = countResultSet.getLong(1);
                        }
                    }
                }
                TableCount tableCount = new TableCount();
                tableCount.setDatabase(databaseName == null ? username : databaseName);
                tableCount.setTable(tableName);
                tableCount.setCount(count);
                tableCounts.add(tableCount);
                log.info("Get count success for " + tableName);
            }
        }
        return tableCounts;
    }
}

MongoDB

package com.seliote.databasecountsnapshot.tool;

import com.mongodb.MongoClient;
import com.mongodb.MongoClientOptions;
import com.mongodb.MongoCredential;
import com.mongodb.ServerAddress;
import com.mongodb.client.MongoCollection;
import com.mongodb.client.MongoDatabase;
import com.seliote.databasecountsnapshot.pojo.MongoConfig;
import com.seliote.databasecountsnapshot.pojo.TableCount;
import lombok.extern.slf4j.Slf4j;
import org.bson.Document;

import java.util.LinkedList;
import java.util.List;

/**
 * MongoDB 相關操作工具類
 *
 * @author LiYangDi
 * @since 2019/12/16
 */
@Slf4j
public class MongoOperate {

    /**
     * 獲取 Mongo 數據庫連接
     *
     * @param mongoConfig MongoDB 連接相關配置
     * @return MongoClient 連接對象
     */
    public static MongoClient getMongoClient(MongoConfig mongoConfig) {
        log.info("Attempt to get DataSource list for: " + mongoConfig);
        ServerAddress serverAddress = new ServerAddress(mongoConfig.getHost(), mongoConfig.getPort());
        MongoCredential mongoCredential = MongoCredential.createCredential(mongoConfig.getUsername(),
                mongoConfig.getDatabase(),
                mongoConfig.getPassword().toCharArray());
        // 空配置即可
        MongoClientOptions mongoClientOptions = MongoClientOptions.builder().build();
        return new MongoClient(serverAddress, mongoCredential, mongoClientOptions);
    }

    /**
     * 關閉 MongoClient 連接
     *
     * @param mongoClient 需要關閉的連接
     */
    public static void closeMongoClient(MongoClient mongoClient) {
        log.info("Close mongo client: " + mongoClient);
        mongoClient.close();
    }

    /**
     * 獲取表中數據量
     *
     * @param mongoClient MongoClient 連接
     * @return 庫中表的數據量
     */
    public static List<TableCount> getTableCounts(MongoClient mongoClient, String databaseName) {
        log.info("Getting MongoDB tables counts");
        List<TableCount> tableCounts = new LinkedList<>();
        // 庫
        MongoDatabase mongoDatabase = mongoClient.getDatabase(databaseName);
        // 遍歷表
        for (String collectionName : mongoDatabase.listCollectionNames()) {
            MongoCollection<Document> mongoCollection = mongoDatabase.getCollection(collectionName);
            TableCount tableCount = new TableCount();
            tableCount.setDatabase(databaseName);
            tableCount.setTable(collectionName);
            tableCount.setCount(mongoCollection.countDocuments());
            tableCounts.add(tableCount);
            log.info("Get count success for " + collectionName);
        }
        return tableCounts;
    }
}