| 提交 | 用户 | 时间 | ||
| e7c126 | 1 | package com.iailab.framework.tenant.core.mq.rocketmq; |
| H | 2 | |
| 3 | import cn.hutool.core.lang.Assert; | |
| 4 | import cn.hutool.core.util.StrUtil; | |
| 5 | import com.iailab.framework.tenant.core.context.TenantContextHolder; | |
| 6 | import org.apache.rocketmq.client.hook.ConsumeMessageContext; | |
| 7 | import org.apache.rocketmq.client.hook.ConsumeMessageHook; | |
| 8 | import org.apache.rocketmq.common.message.MessageExt; | |
| 9 | import org.springframework.messaging.handler.invocation.InvocableHandlerMethod; | |
| 10 | ||
| 11 | import java.util.List; | |
| 12 | ||
| 13 | import static com.iailab.framework.web.core.util.WebFrameworkUtils.HEADER_TENANT_ID; | |
| 14 | ||
| 15 | /** | |
| 16 | * RocketMQ 消息队列的多租户 {@link ConsumeMessageHook} 实现类 | |
| 17 | * | |
| 18 | * Consumer 消费消息时，将消息的 Header 的租户编号，添加到 {@link TenantContextHolder} 中，通过 {@link InvocableHandlerMethod} 实现 | |
| 19 | * | |
| 20 | * @author iailab | |
| 21 | */ | |
| 22 | public class TenantRocketMQConsumeMessageHook implements ConsumeMessageHook { | |
| 23 | ||
| 24 | @Override | |
| 25 | public String hookName() { | |
| 26 | return getClass().getSimpleName(); | |
| 27 | } | |
| 28 | ||
| 29 | @Override | |
| 30 | public void consumeMessageBefore(ConsumeMessageContext context) { | |
| 31 | // 校验，消息必须是单条，不然设置租户可能不正确 | |
| 32 | List<MessageExt> messages = context.getMsgList(); | |
| 33 | Assert.isTrue(messages.size() == 1, "消息条数({})不正确", messages.size()); | |
| 34 | // 设置租户编号 | |
| 35 | String tenantId = messages.get(0).getUserProperty(HEADER_TENANT_ID); | |
| 36 | if (StrUtil.isNotEmpty(tenantId)) { | |
| 37 | TenantContextHolder.setTenantId(Long.parseLong(tenantId)); | |
| 38 | } | |
| 39 | } | |
| 40 | ||
| 41 | @Override | |
| 42 | public void consumeMessageAfter(ConsumeMessageContext context) { | |
| 43 | TenantContextHolder.clear(); | |
| 44 | } | |
| 45 | ||
| 46 | } | |
