はじめに

先日からPostgreSQLのRow Level Securityを使ったマルチテナントアプリケーションを実装することをやっています。

miyohide.hatenablog.com

上の記事では直接curlでテナントIDを指定していましたが、今回はAmazon Cognitoのユーザー属性に追加したカスタム属性をみてRow Level Securityを使ってみることをやってみます。

カスタム属性

カスタム属性はAmazon Cognitoの画面で追加しておきます。頭にcustom:がつくのがポイントです。

ユーザーにそれぞれ値を追加しておきます。

Javaの実装

テナントIDを保存するクラスを作る

Spring Bootは1リクエスト=1スレッドとなるとのことなので、ThreadLocalを使ってテナントIDを保存するクラスを作ります。

public class TenantThreadLocalStorage {
  private static ThreadLocal<String> tenant = new ThreadLocal<>();

  public static String getTenantId() {
    return tenant.get();
  }

  public static void setTenantId(String tenantId) {
    tenant.set(tenantId);
  }
}

リクエストの前処理としてテナントIDを取得する

リクエストの前処理を実装するために、Interceptorを使います。

www.baeldung.com

実装としては以下のような感じになりました。

// リクエストの前処理でcustom:tenant_idを認証情報から取り出し、
// TenantThreadLocalStorageに格納する
@Component
public class RequestInterceptor implements HandlerInterceptor {
  @Override
  public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
    // 認証情報を取り出す
    Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
    if (authentication != null && authentication.getPrincipal() instanceof DefaultOidcUser) {
      DefaultOidcUser user = (DefaultOidcUser) authentication.getPrincipal();
      String tenant_id = user.getClaim("custom:tenant_id");
      TenantThreadLocalStorage.setTenantId(tenant_id);
    }
    return true;
  }

  @Override
  public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) {
    TenantThreadLocalStorage.setTenantId(null);
  }
}

認証情報の取得がかなり苦労しました。SecurityContextHolder.getContext().getAuthentication()で得られたAuthenticationのインスタンスからDefaultOidcUserクラスのインスタンスである場合、custom:tenant_idを取得することがわかったので、上記のような実装になりました。

Interceptorは登録しないといけないのも地味にハマりポイントでした。

// interceptorを登録する
@Configuration
public class WebMvcConfig implements WebMvcConfigurer {
  @Bean
  RequestInterceptor requestInterceptor() {
    return new RequestInterceptor();
  }

  @Override
  public void addInterceptors(InterceptorRegistry registry) {
    registry.addInterceptor(requestInterceptor());
  }
}

取得したテナントIDをPostgreSQLに投げる

これでテナントIDがTenantThreadLocalStorageに保存されているので、データベースとのコネクションを取得する際にTenantThreadLocalStorage.getTenantId()を使って値を取得し、SET app.tenant_id TO 'テナントID'を発行します。

public class TenantAwareDataSource extends DelegatingDataSource {

  public TenantAwareDataSource(DataSource targetDataSource) {
    super(targetDataSource);
  }

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

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

  private void setTenantId(Connection connection) throws SQLException {
    String tenantId;
    try {
      tenantId = TenantThreadLocalStorage.getTenantId();
      if (tenantId == null) {
        tenantId = "-1";
      }
    } catch (ScopeNotActiveException e) {
      tenantId = "-1";
    }
    try (var statement = connection.createStatement()) {
      statement.execute("SET app.tenant_id TO '" + tenantId + "'");
    }
  }

  private Connection getTenantAwareConnectionProxy(Connection connection) {
    return (Connection) Proxy.newProxyInstance(
            ConnectionProxy.class.getClassLoader(),
            new Class[]{ConnectionProxy.class},
            new TenantAwareInvocationHandler(connection)
    );
  }

  static class TenantAwareInvocationHandler implements InvocationHandler {
    private final Connection connection;

    TenantAwareInvocationHandler(Connection connection) {
      this.connection = connection;
    }

    @Override
    public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
      return switch (method.getName()) {
        case "unwrap" -> {
          if (((Class<?>) args[0]).isInstance(proxy)) {
            yield proxy;
          }
          yield method.invoke(connection, args);
        }
        case "isWrapperFor" -> {
          if (((Class<?>) args[0]).isInstance(proxy)) {
            yield true;
          }
          yield method.invoke(connection, args);
        }
        case "close" -> {
          try (var s = connection.createStatement()) {
            s.execute("RESET app.tenant_id");
          }
          yield method.invoke(connection, args);
        }
        case "getTargetConnection" -> connection;
        default -> method.invoke(connection, args);
      };
    }
  }
}

独自のデータソースを作る

上記で作成したTenantAwareDataSourceを返す独自のデータソースを作ります。

@Configuration
public class DatasourceConfig {
  @Bean
  public TenantAwareDataSource datasource(@Value("${myapp.db.user}") String user, @Value("${myapp.db.password}") String password, @Value("${myapp.db.url}") String url) {
    return new TenantAwareDataSource(
            DataSourceBuilder.create().username(user).password(password).url(url).build()
    );
  }
}

実行

あらかじめ以下のようなデータをPostgreSQLのテーブルに保存しておきます。

このとき、テナントIDが1であるtestuser001でログインした場合には、そのデータのみが出力されます。

テナントIDが2であるtestuser002でログインした場合には、そのデータのみが出力されます。

miyohide's blog

Amazon Cognitoカスタム属性を活用したPostgreSQLのRow Level Securityの実装