package com.qianwen.core.tenant.dynamic;

import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import org.aopalliance.intercept.MethodInterceptor;
import org.aopalliance.intercept.MethodInvocation;
import com.qianwen.core.secure.utils.AuthUtil;
import com.qianwen.core.tenant.exception.TenantDataSourceException;
import com.qianwen.core.tool.utils.StringUtil;


public class TenantDataSourceGlobalInterceptor implements MethodInterceptor {
    private TenantDataSourceHolder holder;

    public void setHolder(final TenantDataSourceHolder holder) {
        this.holder = holder;
    }

    public Object invoke(MethodInvocation invocation) throws Throwable {
        String tenantId = AuthUtil.getTenantId();
        try {
            try {
                if (StringUtil.isNotBlank(tenantId)) {
                    this.holder.handleDataSource(tenantId);
                    DynamicDataSourceContextHolder.push(tenantId);
                }
                Object proceed = invocation.proceed();
                if (StringUtil.isNotBlank(tenantId)) {
                    DynamicDataSourceContextHolder.poll();
                }
                return proceed;
            } catch (Exception exception) {
                throw new TenantDataSourceException(exception.getMessage());
            }
        } catch (Throwable th) {
            if (StringUtil.isNotBlank(tenantId)) {
                DynamicDataSourceContextHolder.poll();
            }
            throw th;
        }
    }
}