dashboard-statistics-lambda/src/main/java/com/dashboard/aws/lambda/service/MySQLService.java

package com.dashboard.aws.lambda.service;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.*;
import java.time.LocalDate;
import java.util.*;

public class MySQLService {

    private static final Logger logger = LoggerFactory.getLogger(MySQLService.class);

    private static final String MYSQL_URL = System.getenv("DB_URL"); // 统一实例URL
    private static final String DB_USER = System.getenv("DB_USER");
    private static final String DB_PASSWORD = System.getenv("DB_PASSWORD");


    /**
     * 获取符合条件的企业ID
     */
//    public List<Long> getActiveCompanyIds() throws SQLException {
//        String sql = """
//            SELECT id FROM data_center_new.basic_company
//            WHERE flag != 1
//            AND (parent_id = -1 OR parent_id = 1)
//            """;
//
//        try (Connection conn = DriverManager.getConnection(MYSQL_URL, DB_USER, DB_PASSWORD);
//             PreparedStatement ps = conn.prepareStatement(sql);
//             ResultSet rs = ps.executeQuery()) {
//
//            List<Long> ids = new ArrayList<>();
//            while (rs.next()) {
//                ids.add(rs.getLong("id"));
//            }
//            return ids;
//        }
//    }

    public List<Map<String, Object>> queryOldData(String schema, String table, long startTs, long endTs) throws SQLException {
        String sql = String.format("SELECT * FROM %s.%s WHERE upload_at >= ? AND upload_at < ?", schema, table);
        logger.info("query sql: {}, start: {}, end: {}", sql, startTs, endTs);

        try (Connection conn = DriverManager.getConnection(MYSQL_URL, DB_USER, DB_PASSWORD);
             PreparedStatement ps = conn.prepareStatement(sql)) {

            ps.setLong(1, startTs);
            ps.setLong(2, endTs);

            ResultSet rs = ps.executeQuery();
            ResultSetMetaData meta = rs.getMetaData();
            List<Map<String, Object>> list = new ArrayList<>();

            while (rs.next()) {
                Map<String, Object> row = new LinkedHashMap<>();
                for (int i = 1; i <= meta.getColumnCount(); i++) {
                    row.put(meta.getColumnName(i), rs.getObject(i));
                }
                list.add(row);
            }
            return list;
        }
    }

    public void deleteOldData(String schema, String table, long startTs, long endTs) throws SQLException {
        String sql = String.format("DELETE FROM %s.%s WHERE upload_at >= ? AND upload_at < ?", schema, table);
        logger.info("deleteOldData sql: {}, start: {}, end: {}", sql, startTs, endTs);

        try (Connection conn = DriverManager.getConnection(MYSQL_URL, DB_USER, DB_PASSWORD);
             PreparedStatement ps = conn.prepareStatement(sql)) {

            ps.setLong(1, startTs);
            ps.setLong(2, endTs);
            int deleted = ps.executeUpdate();
            logger.info("[{}.{}] 删除 {} 条记录", schema, table, deleted);
        }
    }

    public int insertRows(String schema, String table, List<Map<String, String>> rows) throws SQLException {
        if (rows == null || rows.isEmpty()) return 0;

        Map<String, String> first = rows.get(0);
        List<String> columns = new ArrayList<>(first.keySet());
        String columnStr = String.join(",", columns);
        String placeholders = String.join(",", Collections.nCopies(columns.size(), "?"));
        String sql = String.format("INSERT INTO %s.%s (%s) VALUES (%s)", schema, table, columnStr, placeholders);
//        logger.info("insert sql: {}", sql);

        int total = 0;
        try (Connection conn = DriverManager.getConnection(MYSQL_URL, DB_USER, DB_PASSWORD);
             PreparedStatement ps = conn.prepareStatement(sql)) {

            for (Map<String, String> row : rows) {
                for (int i = 0; i < columns.size(); i++) {
                    ps.setString(i + 1, row.getOrDefault(columns.get(i), null));
                }
                ps.addBatch();
            }
            int[] results = ps.executeBatch();
            for (int r : results) total += r;
        }
        return total;
    }

    public int deleteBeforeDate(String schema, String table, LocalDate fileDate) {
        // 计算比较值： yyyyMMdd 格式的整数（例如 20251008）
        int cutoff = fileDate.getYear() * 10000 + fileDate.getMonthValue() * 100 + fileDate.getDayOfMonth();

        String sql = String.format(
                "DELETE FROM %s.%s WHERE (COALESCE(date_year,0) * 10000 + COALESCE(date_month,0) * 100 + COALESCE(date_day,0)) < ?",
                schema, table
        );
        logger.info("deleteBeforeDate sql: {}, cutoff:{}", sql, cutoff);

        try (Connection conn = DriverManager.getConnection(MYSQL_URL, DB_USER, DB_PASSWORD);
             PreparedStatement ps = conn.prepareStatement(sql)) {
            ps.setInt(1, cutoff);
            int deleted = ps.executeUpdate();
            logger.info("deleteBeforeDate: schema={}, table={}, cutoff={}, deleted={}", schema, table, cutoff, deleted);
            return deleted;
        } catch (SQLException e) {
            logger.error("deleteBeforeDate failed for {}/{} cutoff={}", schema, table, cutoff, e);
            return 0;
        }
    }

    /**
     * 根据ID列表批量删除数据
     * 使用IN语句一次性删除，提高效率
     */
    public void deleteByIds(String schema, String table, List<Long> ids) throws SQLException {
        if (ids == null || ids.isEmpty()) {
            return;
        }

        // 构建IN子句的占位符字符串：?, ?, ?, ...
        String placeholders = String.join(",", Collections.nCopies(ids.size(), "?"));
        String sql = String.format("DELETE FROM %s.%s WHERE id IN (%s)", schema, table, placeholders);
        logger.info("deleteByIds sql: {}, ids count: {}", sql, ids.size());

        try (Connection conn = DriverManager.getConnection(MYSQL_URL, DB_USER, DB_PASSWORD);
             PreparedStatement ps = conn.prepareStatement(sql)) {

            // 逐个设置ID参数值
            for (int i = 0; i < ids.size(); i++) {
                ps.setLong(i + 1, ids.get(i));
            }
            int deleted = ps.executeUpdate();
            logger.info("[{}.{}] 删除 {} 条记录", schema, table, deleted);
        }
    }

    /**
     * 查询ID大于指定值的前N条数据（按 upload_at + id 升序）
     * 用于分页分批处理数据
     * @param lastId 上次处理的最大ID（不包含），为null时从第一条开始查询
     * @param lastUploadAt 上次处理的最大 upload_at（不包含）
     * @param cutoffUploadAt 截止时间，只处理 upload_at < cutoffUploadAt 的数据
     */
    public List<Map<String, Object>> queryFirstNByIdGreaterThan(String schema, String table, int limit, Long lastId, Long lastUploadAt, Long cutoffUploadAt) throws SQLException {
        String sql;
        if (lastId == null) {
            // 第一次查询
            sql = String.format("SELECT * FROM %s.%s WHERE upload_at < ? ORDER BY upload_at ASC, id ASC LIMIT ?", schema, table);
        } else {
            // 后续查询，从上次之后开始
            sql = String.format("SELECT * FROM %s.%s WHERE (upload_at > ? OR (upload_at = ? AND id > ?)) AND upload_at < ? ORDER BY upload_at ASC, id ASC LIMIT ?", schema, table);
        }
        logger.info("queryFirstNByIdGreaterThan sql: {}, limit: {}, lastId: {}, lastUploadAt: {}, cutoffUploadAt: {}",
                sql, limit, lastId, lastUploadAt, cutoffUploadAt);

        try (Connection conn = DriverManager.getConnection(MYSQL_URL, DB_USER, DB_PASSWORD);
             PreparedStatement ps = conn.prepareStatement(sql)) {

            int paramIndex = 1;
            if (lastId != null) {
                ps.setLong(paramIndex++, lastUploadAt);
                ps.setLong(paramIndex++, lastUploadAt);
                ps.setLong(paramIndex++, lastId);
            }
            ps.setLong(paramIndex++, cutoffUploadAt);
            ps.setInt(paramIndex, limit);

            ResultSet rs = ps.executeQuery();
            ResultSetMetaData meta = rs.getMetaData();
            List<Map<String, Object>> list = new ArrayList<>();

            // 将结果集转换为Map列表
            while (rs.next()) {
                Map<String, Object> row = new LinkedHashMap<>();
                for (int i = 1; i <= meta.getColumnCount(); i++) {
                    row.put(meta.getColumnName(i), rs.getObject(i));
                }
                list.add(row);
            }
            return list;
        }
    }

}