获取 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;
    }
}