Separating Read and Write DataSources in Spring Boot

Introduction

In modern software development, achieving optimal performance and scalability often involves separating read and write operations. This separation allows each operation type to be independently optimized and scaled. In this blog post, we'll explore how to implement separate read and write DataSources in a Spring Boot application using Spring Data JPA and AWS Aurora Serverless PostgreSQL.

Why Separate Read and Write Operations?

Separating read and write operations offers several benefits:

• Performance Optimization: Read-heavy applications can benefit from databases optimized for querying, while write-heavy operations can be handled separately.
• Scalability: Independently scaling read and write operations helps handle different load characteristics.
• Fault Isolation: Issues in one type of operation (e.g., heavy writes) do not affect the other (e.g., read operations).

Setting Up Separate DataSources

We'll use AWS Aurora Serverless PostgreSQL, which provides separate endpoints for read and write operations. Here’s how to configure Spring Boot to use these endpoints.

Step 1: Add Dependencies

Ensure you have the necessary dependencies in your pom.xml:

xml
<dependencies>
    <dependency>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-starter-data-jpa</artifactId>
    </dependency>
    <dependency>
        <groupId>org.postgresql</groupId>
        <artifactId>postgresql</artifactId>
    </dependency>
</dependencies>

Step 2: Configure DataSources

Define the DataSources for reading and writing in your application.yml.

yaml
spring:
  datasource:
    write:
      url: jdbc:postgresql://write-endpoint-url:5432/yourdb
      username: yourusername
      password: yourpassword
      driver-class-name: org.postgresql.Driver
    read:
      url: jdbc:postgresql://read-endpoint-url:5432/yourdb
      username: yourusername
      password: yourpassword
      driver-class-name: org.postgresql.Driver
  jpa:
    hibernate:
      ddl-auto: update
    show-sql: true
    properties:
      hibernate:
        format_sql: true

Step 3: Create Configuration Classes

Create configuration classes to set up the DataSources and routing logic.

ReadWriteDataSourceConfig.java

java
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.boot.jdbc.DataSourceBuilder;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

import javax.sql.DataSource;
import java.util.HashMap;
import java.util.Map;

@Configuration
public class ReadWriteDataSourceConfig {

    @Bean(name = "writeDataSource")
    @ConfigurationProperties(prefix = "spring.datasource.write")
    public DataSource writeDataSource() {
        return DataSourceBuilder.create().build();
    }

    @Bean(name = "readDataSource")
    @ConfigurationProperties(prefix = "spring.datasource.read")
    public DataSource readDataSource() {
        return DataSourceBuilder.create().build();
    }

    @Bean(name = "dataSource")
    public DataSource dataSource(@Qualifier("writeDataSource") DataSource writeDataSource,
                                 @Qualifier("readDataSource") DataSource readDataSource) {
        AbstractRoutingDataSource routingDataSource = new ReplicationRoutingDataSource();
        Map<Object, Object> dataSourceMap = new HashMap<>();
        dataSourceMap.put(DataSourceType.WRITE, writeDataSource);
        dataSourceMap.put(DataSourceType.READ, readDataSource);

        routingDataSource.setDefaultTargetDataSource(writeDataSource);
        routingDataSource.setTargetDataSources(dataSourceMap);
        return routingDataSource;
    }
}

ReplicationRoutingDataSource.java

java
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;

public class ReplicationRoutingDataSource extends AbstractRoutingDataSource {

    @Override
    protected Object determineCurrentLookupKey() {
        return DataSourceContextHolder.getDataSourceType();
    }
}

Step 4: Implement DataSource Routing Logic

DataSourceType.java

java
public enum DataSourceType {
    READ, WRITE
}

DataSourceContextHolder.java

java
public class DataSourceContextHolder {
    private static final ThreadLocal<DataSourceType> CONTEXT_HOLDER = new ThreadLocal<>();

    public static void setDataSourceType(DataSourceType dataSourceType) {
        CONTEXT_HOLDER.set(dataSourceType);
    }

    public static DataSourceType getDataSourceType() {
        return CONTEXT_HOLDER.get();
    }

    public static void clearDataSourceType() {
        CONTEXT_HOLDER.remove();
    }
}

Step 5: Create AOP Aspect for Routing

Use Aspect-Oriented Programming (AOP) to route read and write operations to the correct DataSource.

DataSourceAspect.java

java
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Before;
import org.springframework.stereotype.Component;

@Aspect
@Component
public class DataSourceAspect {

    @Before("execution(* com.yourpackage.repository..*.save*(..)) || execution(* com.yourpackage.repository..*.insert*(..)) || execution(* com.yourpackage.repository..*.update*(..)) || execution(* com.yourpackage.repository..*.delete*(..))")
    public void setWriteDataSourceType() {
        DataSourceContextHolder.setDataSourceType(DataSourceType.WRITE);
    }

    @Before("execution(* com.yourpackage.repository..*.find*(..)) || execution(* com.yourpackage.repository..*.get*(..)) || execution(* com.yourpackage.repository..*.read*(..))")
    public void setReadDataSourceType() {
        DataSourceContextHolder.setDataSourceType(DataSourceType.READ);
    }
}

Step 6: Configure JPA EntityManager

Configure the LocalContainerEntityManagerFactoryBean to use the custom DataSource.

JpaConfig.java

java
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.boot.orm.jpa.EntityManagerFactoryBuilder;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.data.jpa.repository.config.EnableJpaRepositories;
import org.springframework.orm.jpa.JpaTransactionManager;
import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean;
import org.springframework.transaction.PlatformTransactionManager;

import javax.sql.DataSource;

@Configuration
@EnableJpaRepositories(
        basePackages = "com.yourpackage.repository",
        entityManagerFactoryRef = "entityManagerFactory",
        transactionManagerRef = "transactionManager"
)
public class JpaConfig {

    @Bean(name = "entityManagerFactory")
    public LocalContainerEntityManagerFactoryBean entityManagerFactory(
            EntityManagerFactoryBuilder builder,
            @Qualifier("dataSource") DataSource dataSource) {
        return builder
                .dataSource(dataSource)
                .packages("com.yourpackage.domain")
                .persistenceUnit("default")
                .build();
    }

    @Bean(name = "transactionManager")
    public PlatformTransactionManager transactionManager(
            @Qualifier("entityManagerFactory") LocalContainerEntityManagerFactoryBean entityManagerFactory) {
        return new JpaTransactionManager(entityManagerFactory.getObject());
    }
}

Conclusion

By following these steps, you can effectively separate read and write DataSources in your Spring Boot application. This approach provides better performance, scalability, and fault isolation. The use of AOP ensures that read and write operations are correctly routed, maintaining a clean separation between them. This setup is particularly beneficial for applications with heavy read and write operations, allowing each to be optimized and scaled independently.

Implementing such a pattern can lead to significant improvements in the efficiency and reliability of your application, making it better suited to handle varying loads and operational requirements.