Dynamic database connection switching and application multitenancy with database routing

by Emil Korladinov

Introduction

Projects with  requirements to dynamically switch connections at runtime  and exchange data with different datasources along with multi-tenat capabilities are very common. 

For a project I worked on I compiled a solution using Spring Framework, Spring Data Jpa, Hibernate and Aspect Oriented Programming  with Spring (SpringAOP). 

Here  I'd like to share my solution with all of you.

Implementation

Basic idea is to wrap unit of work into a transaction and to exexute it against preconfigured databases(datasources). Transactions are demarcated using standard Spring annotation-based transaction management. Since annotations are used to mark code which will use database routing a new annotation is introduced: @DatabaseRouting which is meta annotated with Spring @Transactional. Information for the currently used tenant id is kept in a ThreadLocal variable on a per-thread basis. On the database side a special DataSource is used - DatabaseRoutingDatasource which extends Spring's AbstractRoutingDataSource. 

Working horse of the solution is aspect with an @Around pointcut. Second annotation, applied on method parameter is used to convey information for chosen datasource down to Spring transaction.

Annotations

Two annotations are used:

@DatabaseRouting is method level annotation marking method which uses database routing to dynamically select database connection and exchange data. Methods that have this anotation are wrapped in transaction. This annotation is matched by an around aspect which sets current connection discriminator (tenantId), proceeds with methd execution and resets connection to default.

This annotation is defined as follows:

	@Retention(RetentionPolicy.RUNTIME)
	@Target(ElementType.METHOD)
	@Transactional(propagation = Propagation.REQUIRES_NEW)
	public @interface DatabaseRouting {
	}

view raw DatabaseRouting.java hosted with ❤ by GitHub

@TenantId is parameter level annotation, used as a vehicle for required connection discriminator (tenantId) to the processing aspect. Parameter with this annotation is not used for business logic in the method body, but removing it from the method signature with break processing and method will be executed on the default connection.

	@Retention(RetentionPolicy.RUNTIME)
	@Target(ElementType.PARAMETER)
	public @interface TenantId {
	}

view raw TenantId.java hosted with ❤ by GitHub

ThreadLocals

TenantIdHolder is standard Java ThreadLocal containing current value of tenantId.

	public class TenantIdHolder {
	private static final ThreadLocal<String> holder = new ThreadLocal<>();
	public static String get() {
	return holder.get();
	}
	public static void set(String tenantId) {
	if (null == tenantId) {
	throw new NullPointerException();
	}
	holder.set(tenantId);
	}
	public static void clear() {
	holder.remove();
	}
	}

view raw TenantIdHolder.java hosted with ❤ by GitHub

Datasources

RoutingDataSource Extends AbstractRoutingDataSource from spring, overriding method for key lookup (tenantId), linking it with TenantIdHolder.

Aspects

DatabaseRoutingAspect

This is the working horse of the solution. Purpose of this aspect is twofold - First to extract tenantId from the method arguments through annotated parameter on the method signature; and second execute business logic of the method and reset connection to default.

 

	@Component
	@Aspect
	@Order(200)
	public class DatebaseRoutingAspect {
	@Autowired
	private Map dataSources;
	@Pointcut("@annotation(DatabaseRouting))")
	public void databaseRouting() {
	}
	@Pointcut("execution(* *(..))")
	public void atExecution() {
	}
	@Around("@annotation(databaseRouting)")
	public Object switchDatabase(ProceedingJoinPoint pjp, DatabaseRouting databaseRouting) throws Throwable {
	String tenantId = retrieveTenantId(pjp);
	try {
	if (null != tenantId) {
	if (!dataSources.containsKey(tenantId)) {
	throw new IllegalArgumentException("Unknown database key: " + tenantId);
	}
	TenantIdHolder.set(tenantId);
	}
	return pjp.proceed();
	} finally {
	TenantIdHolder.clear();
	}
	}
	private String retrieveTenantId(ProceedingJoinPoint pjp) {
	MethodSignature signature = (MethodSignature) pjp.getSignature();
	Annotation[][] parameterAnnotations = signature.getMethod().getParameterAnnotations();
	Object[] args = pjp.getArgs();
	int idx = 0;
	Object tenantId = null;
	for (Annotation[] annotation : parameterAnnotations) {
	tenantId = args[idx++];
	for (Annotation annot : annotation) {
	if (TenantId.class.equals(annot.annotationType())) {
	break;
	}
	}
	}
	return (String) tenantId;
	}
	}

view raw DatebaseRoutingAspect.java hosted with ❤ by GitHub

Sources can be found on GitHub