【SpringBoot】動的にDBを切り替える方法

はじめに

エキサイト株式会社 バックエンドエンジニアの山縣です。

既存サービスのリビルドをするにあたって、SpringBootで動的にDBを切り替える必要がありました。 これを実現するためにやったことについて本記事で紹介します。

動作環境

  • SpringBoot 2.6.2
  • Java 17
  • MyBatis Spring Boot Starter: 2.1.3
  • SQL Server 2016

複数のDB

現在私が携わっているサービスでは「ユーザーIDのプレフィックス1文字を見てDBの向き先を変更する」というような負荷分散が取り入れられています。 下記図のように、例えばIDが101234567であればDB1を参照し、IDが987654321であればDB9を参照するイメージです。 そのため、アプリケーション側ではIDをもとに動的にDBの向き先を変更する必要がありました。

DataSourceの設定

まずはDBの接続情報をapplication.ymlに記述します。 今回は10台のDBを使用して負荷分散することを想定しているので、10台分の設定が必要です。 例として2台分記述していますが、残り8台のDBの設定情報も同じように記述します。

spring:
  datasource:
    db01:
      url: jdbc:sqlserver://127.0.0.1:1433;databaseName=db01
      username: user1
      password: pass1
      driverClassName: com.microsoft.sqlserver.jdbc.SQLServerDriver
    db02:
      url: jdbc:sqlserver://127.0.0.1:1433;databaseName=db02
      username: user2
      password: pass2
      driverClassName: com.microsoft.sqlserver.jdbc.SQLServerDriver

次に@ConfigurationPropertiesを使用してDBの接続情報をDataSourceとして扱えるようにします。 上記同様に残り8台分の設定情報は省略しています。

package com.example.db.config;

import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

@Configuration
public class DataSourcePropertiesConfig {
    public static final String DATA_SOURCE_PROPERTIES_DB01 = "dataSourcePropertiesDb01";
    public static final String DATA_SOURCE_PROPERTIES_DB02 = "dataSourcePropertiesDb02";

    public static final String DATA_SOURCE_DB01 = "dataSourceDb01";
    public static final String DATA_SOURCE_DB02 = "dataSourceDb02";

    @Bean(name = {DATA_SOURCE_PROPERTIES_DB01})
    @ConfigurationProperties(prefix = "spring.datasource.db01")
    public DataSourceProperties dataSourcePropertiesDb01() {
        return new DataSourceProperties();
    }

    @Bean(name = {DATA_SOURCE_PROPERTIES_DB02})
    @ConfigurationProperties(prefix = "spring.datasource.db02")
    public DataSourceProperties dataSourcePropertiesDb02() {
        return new DataSourceProperties();
    }

    @Bean(name = {DATA_SOURCE_DB01})
    public DataSource dataSourceDb01(@Qualifier(DATA_SOURCE_PROPERTIES_DB01) DataSourceProperties properties) {
        return properties
                .initializeDataSourceBuilder()
                .build();
    }

    @Bean(name = {DATA_SOURCE_DB02})
    public DataSource dataSourceDb02(@Qualifier(DATA_SOURCE_PROPERTIES_DB02) DataSourceProperties properties) {
        return properties
                .initializeDataSourceBuilder()
                .build();
    }
}

DBの接続先を定義したEnumの作成

DBの接続先を定義したEnumを作成します。 IDのプレフィックス1文字とDBの接続先を対応づけています。

また、IDからDBの接続先のEnumを取得できるようにしており、 DatabaseSchemeType.getSchemeType("312345678")DB03と取得することができるように実装しました。

package com.example.sqlserver.config;

import lombok.Getter;
import org.apache.commons.lang3.StringUtils;

import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public enum DatabaseSchemeType {
    DB00("0"),
    DB01("1"),
    DB02("2"),
    DB03("3"),
    DB04("4"),
    DB05("5"),
    DB06("6"),
    DB07("7"),
    DB08("8"),
    DB09("9");

    @Getter
    private final String prefix;

    DatabaseSchemeType(String prefix) {
        this.prefix = prefix;
    }

    private static final Map<String, DatabaseSchemeType> DATABASE_SCHEME_TYPE_MAP = Stream
            .of(DatabaseSchemeType.values())
            .collect(Collectors.toMap(
                    e -> e.getPrefix(),
                    e -> e
            ));

    private static final Pattern ID_PATTERN = Pattern.compile("^([0-9])[0-9]{8}$");

    /**
     * DBスキーマを習得
     *
     * @param id
     * @return DBスキーマ
     */
    public static DatabaseSchemeType getSchemeType(String id) {
        final Matcher matcher = ID_PATTERN.matcher(id);
        if (!matcher.find()) {
            throw new RuntimeException("unexpected id pattern");
        }

        final String prefix = matcher.group(1);
        return DATABASE_SCHEME_TYPE_MAP.get(prefix);
    }
}

DataSourceContextHolderの作成

ThreadLocalを使用してDBの接続先を扱います。 DBに接続する箇所でThreadLocalDatabaseSchemeTypeをセットし、実行が終わったあとにクリアされるようにします。 使い方は後ほど説明します。

package com.example.db.config;

public class DatabaseContextHolder {
    private static final ThreadLocal<DatabaseSchemeType> contextHolder = new ThreadLocal<>();

    public static void setDatabaseSchemeType(DatabaseSchemeType type) {
        contextHolder.set(type);
    }

    public static DatabaseSchemeType getDatabaseSchemeType() {
        return contextHolder.get();
    }

    public static void clear() {
        contextHolder.remove();
    }
}

AbstractRoutingDataSourceを実装する

AbstractRoutingDataSourcedetermineCurrentLookupKeyメソッドを実装することでDataSourceを動的に切り替えることができるようになります。 DatabaseContextHolderからDatabaseSchemeTypeを取得して、それと対応するDataSourceのBean名を返しています。

package com.example.sqlserver.config;

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

import java.util.Map;
import java.util.Optional;

import static com.example.sqlserver.config.DataSourceConfig.DATA_SOURCE_DB00;
import static com.example.sqlserver.config.DataSourceConfig.DATA_SOURCE_DB01;
import static com.example.sqlserver.config.DataSourceConfig.DATA_SOURCE_DB02;
import static com.example.sqlserver.config.DataSourceConfig.DATA_SOURCE_DB03;
import static com.example.sqlserver.config.DataSourceConfig.DATA_SOURCE_DB04;
import static com.example.sqlserver.config.DataSourceConfig.DATA_SOURCE_DB05;
import static com.example.sqlserver.config.DataSourceConfig.DATA_SOURCE_DB06;
import static com.example.sqlserver.config.DataSourceConfig.DATA_SOURCE_DB07;
import static com.example.sqlserver.config.DataSourceConfig.DATA_SOURCE_DB08;
import static com.example.sqlserver.config.DataSourceConfig.DATA_SOURCE_DB09;

public class RoutingDataSourceResolver extends AbstractRoutingDataSource {
    /**
     * key: DBスキーマ, value: DataSourceのBean名 のマップ
     */
    private static final Map<DatabaseSchemeType, String> DATA_SOURCE_BEAN_NAME_MAP = Map.of(
            DatabaseSchemeType.DB00, DATA_SOURCE_DB00,
            DatabaseSchemeType.DB01, DATA_SOURCE_DB01,
            DatabaseSchemeType.DB02, DATA_SOURCE_DB02,
            DatabaseSchemeType.DB03, DATA_SOURCE_DB03,
            DatabaseSchemeType.DB04, DATA_SOURCE_DB04,
            DatabaseSchemeType.DB05, DATA_SOURCE_DB05,
            DatabaseSchemeType.DB06, DATA_SOURCE_DB06,
            DatabaseSchemeType.DB07, DATA_SOURCE_DB07,
            DatabaseSchemeType.DB08, DATA_SOURCE_DB08,
            DatabaseSchemeType.DB09, DATA_SOURCE_DB09,
    );

    @Override
    protected Object determineCurrentLookupKey() {
        final DatabaseSchemeType type =DatabaseContextHolder.getDatabaseSchemeType();
        return DATA_SOURCE_BEAN_NAME_MAP.get(type);
    }
}

SqlSessionFactoryを定義する

最後に、作成したRoutingDataSourceResolverを使用してSqlSessionFactoryのBeanを定義します。 これでSpringBootで動的にDBを切り替えるための準備が終わりました。

package com.example.sqlserver.config;

import org.apache.ibatis.session.SqlSessionFactory;
import org.mybatis.spring.SqlSessionFactoryBean;
import org.mybatis.spring.SqlSessionTemplate;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import static com.example.sqlserver.config.DataSourceConfig.ROUTING_DATA_SOURCE_RESOLVER;
import static com.example.sqlserver.config.RoutingDataSourceResolver.DATA_SOURCE_BEAN_NAME_MAP;

@Configuration
public class SqlSessionFactoryConfig {
    public static final String SQL_SESSION_FACTORY = "sqlSessionFactory";

    @Bean(name = {SQL_SESSION_FACTORY})
    public SqlSessionFactory sqlSessionFactory(@Qualifier(ROUTING_DATA_SOURCE_RESOLVER) RoutingDataSourceResolver resolver) throws Exception {
        SqlSessionFactoryBean sqlSessionFactoryBean = new SqlSessionFactoryBean();
        sqlSessionFactoryBean.setDataSource(resolver);
        SqlSessionFactory sqlSessionFactory = sqlSessionFactoryBean.getObject();

        // MyBatisの設定を記述

        return sqlSessionFactory;
    }
}

実際に実行する

下記手順で実行すると動的にDBを切り替えることができるようになります。

  1. IDから接続先のEnumを取得する
  2. DatabaseContextHolderに接続先のEnumをセットする
  3. データ取得やデータ追加などのクエリーを実行する
  4. DatabaseContextHolderにセットしたEnumを削除する
package com.examle.repository;

import org.springframework.stereotype.Repository;

@Repository
@RequiredArgsConstructor
public class ItemRepositoryImpl implements ItemRepository {

    @Override
    public Item findItemById(String id) {
        final DatabaseSchemeType type = DatabaseSchemeType.getSchemeType(id);
        DatabaseContextHolder.setDatabaseSchemeType(type);

        /* データ取得の処理 */

        DatabaseContextHolder.clear();
    }
}

おわりに

本記事では、SpringBootで動的にDBを切り替える方法について紹介しました。 今回はクエリー実行箇所で毎回DatabaseSchemeTypeをセットする方法について書きましたがこれは非常に面倒です。 そのためAOPを使用して自動で切り替えられるようにすることで手間を省くのがよいと思います。 これについてはまたの機会にまとめる予定です。 最後まで読んでいただき、ありがとうございました!

参考