はじめに

先日、マルチテナントなアプリケーションの実装のためにPostgreSQLのRow Level Securityの検証結果を記しました。

miyohide.hatenablog.com

今回はRow Level Securityを設定したPostgreSQLのテーブルに対してSpring Bootアプリケーションからアクセスすることを実装してみます。

参考

実装方法は色々とあるようなのですが、今回は以下のブログの内容を参考に実装しました。

www.wirequery.io

実装

実装の肝となるのは、app.current_tenant_idの設定です。今回は、リクエストごとにクライアントからテナントIDが送られてくると想定して、@RequestScopeを使うことにします。以下のQiitaの記事を参考にしました。

qiita.com

これをもとに、以下のクラスを作成します。

@Component
@RequestScope
public class TenantRequestContext {
  private Integer tenantId;

  // getter/setterは省略
}

HTTPリクエストからテナントIDに該当するものを抜き出し、TenantRequestContextのインスタンスに値を設定する処理はOncePerRequestFilterを拡張したクラスで実装します。

@Component
public class TenantAwareFilter extends OncePerRequestFilter {
  private final TenantRequestContext tenantRequestContext;

  public TenantAwareFilter(TenantRequestContext tenantRequestContext) {
    this.tenantRequestContext = tenantRequestContext;
  }

  @Override
  protected void doFilterInternal(HttpServletRequest request, HttpServletResponse response, FilterChain filterChain) throws ServletException, IOException {
    var tenantId = request.getIntHeader("Tenant-Id");
    tenantRequestContext.setTenantId(tenantId);
    // テナントIDが取得できない場合は400エラーを呼び出し元に返す
    if (tenantId == -1) {
      response.setStatus(HttpServletResponse.SC_BAD_REQUEST);
      return;
    }
    filterChain.doFilter(request, response);
  }
}

リクエスト元からテナントID（ここではTenant-Id）がヘッダとして指定されていなかった場合、400エラーを呼び出し元に返し以降の処理は行わないようにしています。

次に、DelegatingDataSourceを拡張したクラスを作り、app.current_tenant_idを設定する処理を実装します。

public class TenantAwareDataSource extends DelegatingDataSource {
  private final TenantRequestContext tenantRequestContext;

  public TenantAwareDataSource(TenantRequestContext tenantRequestContext, DataSource dataSource) {
    super(dataSource);
    this.tenantRequestContext = tenantRequestContext;
  }

  @Override
  public Connection getConnection() throws SQLException {
    var connection = Objects.requireNonNull(getTargetDataSource().getConnection());
    setTenantId(connection);
    return getTenantAwareConnectionProxy(connection);
  }

  @Override
  public Connection getConnection(String username, String password) throws SQLException {
    var connection = Objects.requireNonNull(getTargetDataSource()).getConnection();
    setTenantId(connection);
    return getTenantAwareConnectionProxy(connection);
  }

  private void setTenantId(Connection connection) throws SQLException {
    Integer tenantId;
    try {
      tenantId = tenantRequestContext.getTenantId();
      if (tenantId == null) {
        tenantId = -1;
      }
    } catch (ScopeNotActiveException e) {
      tenantId = -1;
    }
    // Try with resourcesの使用
    try (var statement = connection.createStatement()) {
      statement.execute("SET app.current_tenant_id TO '" + tenantId + "'");
    }
  }
  // 以下省略

ちょっとハマったのは、setTenantIdメソッドの実装で、Spring Bootアプリの起動時にScopeNotActiveExceptionが発生するという点です。Spring Bootアプリの起動時にコネクションプールやDatabase Migrationのための接続のためにリクエストスレッド以外でtenantRequestContextにアクセスしてしまうためScopeNotActiveExceptionが発生したというのがその理由でした。上記の実装のようにtry-catchで囲ってあげれば解決しました。

データベースに接続するための独自のデータソースも必要となります。以下の資料を参照に。

spring.pleiades.io

上で実装したTenantAwareDataSourceのインスタンスを返すのが目的です。

@Configuration
public class DatasourceConfig {
  @Bean
  public TenantAwareDataSource dataSource(
          TenantRequestContext tenantRequestContext,
          @Value("${myapp.db.user}") String user,
          @Value("${myapp.db.password}") String password,
          @Value("${myapp.db.url}") String url
  ) {
    return new TenantAwareDataSource(
            tenantRequestContext,
            DataSourceBuilder.create()
                    .username(user)
                    .password(password)
                    .driverClassName("org.postgresql.Driver")
                    .url(url)
                    .build()
    );
  }
}

実行

以下のような形でテーブルデータがあるとします。

postgres=# SELECT * FROM customers;
 id | tenant_id |       name        | address 
----+-----------+-------------------+---------
  1 |         1 | tenant 1 customer | Unknown
  2 |         2 | tenant 2 customer | Unknown
(2 rows)

postgres=#

リクエストにテナントIDを付与してアクセスすると、無事、テナントIDに対応したものだけが返ってきます。

% curl localhost:8080/customers -H 'Tenant-Id: 1'
[{"id":1,"name":"tenant 1 customer","address":"Unknown"}]%                                                                        
% curl localhost:8080/customers -H 'Tenant-Id: 2'
[{"id":2,"name":"tenant 2 customer","address":"Unknown"}]%                                                                        
% curl localhost:8080/customers                  
%