小米训练营-第六次作业

[!NOTE]

请设计基于Spring boot、Mybaits实现一个简化版的分表查询功能，要点解释：

不可采用sharding-jdbc实现，可以基于 Mybaits 拦截器实现、也可以基于JDBC实现

共两个库db0、db1，其中每个库两张表user_0、user_1 累计4张表

实现根据用户id的查询和插入逻辑

⭐实现基于用户email的模糊查询

1. Mybatis 的 TableRoutingInterceptor 实现

根据以上三个数值，对userId取模后的值进行分库分表。

分库逻辑：

分表逻辑：

@Override
public Object intercept(Invocation invocation) throws Throwable {
    // 获取调用参数
    Object[] args = invocation.getArgs();
    MappedStatement ms = (MappedStatement) args[0];
    Object parameter = args[1];

    String methodId = ms.getId(); // 获取 MyBatis 方法 ID（如 mapper 接口全限定名 + 方法名）
    String originalSql = ms.getBoundSql(parameter).getSql(); // 获取原始 SQL 语句

    try {
        // 根据方法 ID 分发到不同的处理逻辑
        if (methodId.endsWith("insertUser")) {
            return handleInsert(invocation, originalSql, parameter);
        } else if (methodId.endsWith("getUserById")) {
            return handleIdQuery(invocation, originalSql, parameter);
        } else if (methodId.endsWith("getUserByEmail")) {
            return handleEmailQuery(invocation, originalSql, parameter);
        }
        return invocation.proceed(); // 如果没有匹配的方法，直接执行原 SQL
    } finally {
        DataSourceConfig.DynamicDataSource.clear(); // 清除动态数据源设置
    }
}

intercept方法会根据调用的 MyBatis 方法决定如何修改 SQL 或切换数据源。

private Object handleInsert(Invocation invocation, String sql, Object parameter) throws Throwable {
    User user = (User) parameter;
    if (user.getId() == null) {
        throw new IllegalArgumentException("User ID cannot be null");
    }

    int shard = calculateShard(user.getId());// 计算userId % TOTAL_SHARDS 其中 TOTAL_SHARDS = 2
    String newSql = replaceTableName(sql, shard % TABLE_COUNT_PER_DB); //TABLE_COUNT_PER_DB = 2 * 2
    setDataSource(shard / TABLE_COUNT_PER_DB);

    return executeWithNewSql(invocation, newSql);
}

handleInsert方法处理用户插入操作，计算分片位置并重写表名，设置对应的数据源，并执行新的 SQL。

private Object handleIdQuery(Invocation invocation, String sql, Object parameter) throws Throwable {
    Long userId = (Long) parameter;
    int shard = calculateShard(userId);
    String newSql = replaceTableName(sql, shard % TABLE_COUNT_PER_DB);
    setDataSource(shard / TABLE_COUNT_PER_DB);

    return executeWithNewSql(invocation, newSql);
}

handleIdQuery方法处理基于 ID 的精确查询，根据用户 ID 计算分片索引，替换表名并选择对应的数据源，执行新生成的 SQL。

private Object handleEmailQuery(Invocation invocation, String sql, Object parameter) throws Throwable {
    List<User> results = new ArrayList<>();
    Executor executor = (Executor) invocation.getTarget();

    for (int dbIndex = 0; dbIndex < DB_COUNT; dbIndex++) {
        for (int tableIndex = 0; tableIndex < TABLE_COUNT_PER_DB; tableIndex++) {
            try {
                setDataSource(dbIndex);
                String newSql = replaceTableName(sql, tableIndex);
                Object result = executeWithNewSql(invocation, newSql, executor);
                if (result instanceof List) {
                    results.addAll((List<User>) result);
                }
            } catch (Exception e) {
                System.err.println("Error querying db" + dbIndex + ".user_" + tableIndex + ": " + e.getMessage());
            }
        }
    }
    return results;
}

handleEmailQuery处理基于 Email 的模糊查询。因为无法确定数据所在分片，需对所有分片进行搜索。遍历所有数据库和表，逐个查询并将结果合并返回。其中executeWithNewSql用来构造并执行的SQL。

private MappedStatement copyMappedStatement(MappedStatement ms, BoundSql boundSql) {
    MappedStatement.Builder builder = new MappedStatement.Builder(
        ms.getConfiguration(),
        ms.getId(),
        new SqlSource() {
            @Override
            public BoundSql getBoundSql(Object parameterObject) {
                return boundSql;
            }
        },
        ms.getSqlCommandType()
    );

    builder.resource(ms.getResource())
        .fetchSize(ms.getFetchSize())
        .statementType(ms.getStatementType())
        .keyGenerator(ms.getKeyGenerator())
        .keyProperty(ms.getKeyProperties() == null ? null : String.join(",", ms.getKeyProperties()))
        .timeout(ms.getTimeout())
        .parameterMap(ms.getParameterMap())
        .resultMaps(ms.getResultMaps())
        .resultSetType(ms.getResultSetType())
        .cache(ms.getCache())
        .flushCacheRequired(ms.isFlushCacheRequired())
        .useCache(ms.isUseCache());

    return builder.build();
}

以上查询ID、插入用户、查询Email方法都会调用replaceTableName方法，通过修改 SQL 中的表名来实现：

private String replaceTableName(String sql, int tableSuffix) {
    return sql.replaceAll("user(_\\d+)?", "user_" + tableSuffix);//根据UserId % 2 得出
}

2. UserServiceImpl 实现类

package cn.yomigaeri.xiaomi.service.impl;

import cn.yomigaeri.xiaomi.entity.po.User;
import cn.yomigaeri.xiaomi.mapper.UserMapper;
import cn.yomigaeri.xiaomi.service.UserService;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.List;

/**
 * @ClassName: UserServiceImpl
 * @Package: cn.yomigaeri.xiaomi.service.impl
 * @Description:
 * @Author Yomigaeri
 * @Create 2025/5/29 17:00
 * @Version 1.0
 */
@Service
public class UserServiceImpl implements UserService {
    @Autowired
    private UserMapper userMapper;

    public User getUserById(Long id) {
        return userMapper.getUserById(id);
    }

    public void insertUser(User user) {
        if (user.getId() == null) {
            user.setId(generateId());
        }
        userMapper.insertUser(user);
    }

    public List<User> getUserByEmail(String email) {
        return userMapper.getUsersByEmail("%" + email + "%");
    }

    private Long generateId() {
        // 分布式ID生成算法（简单化模拟）
        return System.currentTimeMillis() * 1000 + (long) (Math.random() * 1000);
    }
}