i have a multi tenancy application written with spring boot, keycloak, flyway and postgres, i have been finding it difficult to switch from one schema to another, note i have done a lot of research and what i got is limited to executing jpa query methods and not custom quey methods. Here is my code structure:
This is my TenantContext class

public final class TenantContext {
    private static final String LOGGER_TENANT_ID = "tenant_id";
    public static final String DEFAULT_TENANT = "public";
    private static final ThreadLocal<String> CURRENT_TENANT = new ThreadLocal<>();

    private TenantContext() {}

    public static String getCurrentTenant() {
        String tenant = CURRENT_TENANT.get();
        if (tenant != null) {
            return tenant;
        }
        return DEFAULT_TENANT;
    }

    public static void setCurrentTenant(String tenant) {
        MDC.put(LOGGER_TENANT_ID, tenant);
        CURRENT_TENANT.set(tenant);
    }

    public static void clear() {
        MDC.clear();
        CURRENT_TENANT.remove();
    }
}

And my context filter class:

@Component
public class TenantContextFilter implements Filter {
    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {
        try {
            if (TenantContext.getCurrentTenant().equals(TenantContext.DEFAULT_TENANT)) {
                Authentication authentication = SecurityContextHolder.getContext().getAuthentication();
                if (authentication != null && authentication.getPrincipal() instanceof OidcUser oidcUser) {
                    String tenant = (String) oidcUser.getAttributes().get("azp");
                    TenantContext.setCurrentTenant(Constants.PLATFORM_PREFIX + tenant.toLowerCase());
                } else if ((authentication instanceof JwtAuthenticationToken jwtAuthenticationToken)) {
                    Jwt jwt = jwtAuthenticationToken.getToken();
                    String tenant = jwt.getClaimAsString("azp");
                    TenantContext.setCurrentTenant(Constants.PLATFORM_PREFIX + tenant.toLowerCase());
                } else {
                    HttpServletRequest httpRequest = (HttpServletRequest) request;
                    String requestURI = httpRequest.getRequestURI();
                    if (requestURI.contains("/api/v1/admin")) {
                        TenantContext.setCurrentTenant("public");
                    } else {
                        ((HttpServletResponse) response).sendError(HttpServletResponse.SC_UNAUTHORIZED, "Unauthorized");
                        return;
                    }
                }
            }
            chain.doFilter(request, response);
        } finally {
            TenantContext.clear();
        }
    }

    @Override
    public void init(FilterConfig filterConfig) throws ServletException {
        // Initialization code, if needed
    }

    @Override
    public void destroy() {
        // Cleanup code, if needed
    }
}

My HibernateConfig class:

@Configuration
public class HibernateConfig {

    @Autowired
    private DataSource dataSource;

    @Autowired
    private MultiTenantConnectionProviderImpl multiTenantConnectionProvider;

    @Autowired
    private CurrentTenantIdentifierResolverImpl currentTenantIdentifierResolver;

    @Bean
    public LocalContainerEntityManagerFactoryBean entityManagerFactory() {
        LocalContainerEntityManagerFactoryBean em = new LocalContainerEntityManagerFactoryBean();
        em.setDataSource(dataSource);
        em.setPackagesToScan("com.mahshellsoft.bims");
        HibernateJpaVendorAdapter vendorAdapter = new HibernateJpaVendorAdapter();
        em.setJpaVendorAdapter(vendorAdapter);
        Map<String, Object> properties = new HashMap<>();
        properties.put("hibernate.multiTenancy", "SCHEMA");
        properties.put("hibernate.multi_tenant_connection_provider", multiTenantConnectionProvider);
        properties.put("hibernate.tenant_identifier_resolver", currentTenantIdentifierResolver);
        em.setJpaPropertyMap(properties);
        return em;
    }
}

My MultiTenantConnectionProviderImpl class:

@Component
public class MultiTenantConnectionProviderImpl implements MultiTenantConnectionProvider {
    @Autowired
    private DataSource dataSource;

    @Override
    public Connection getAnyConnection() throws SQLException {
        return dataSource.getConnection();
    }

    @Override
    public void releaseAnyConnection(Connection connection) throws SQLException {
        connection.close();
    }

    @Override
    public Connection getConnection(Object tenantIdentifier) throws SQLException {
        final Connection connection = getAnyConnection();
        try {
            MDC.put("schema", tenantIdentifier.toString());
            connection.setSchema((String) tenantIdentifier);
            connection.createStatement().execute("SET search_path TO " + tenantIdentifier);
        } catch (SQLException e) {
            throw new HibernateException("Could not set schema to " + tenantIdentifier, e);
        }
        return connection;
    }

    @Override
    public void releaseConnection(Object tenantIdentifier, Connection connection) throws SQLException {
        try {
            connection.setSchema("public");
            connection.createStatement().execute("SET search_path TO public");
        } catch (SQLException e) {
            throw new HibernateException("Could not reset schema to public", e);
        }
        connection.close();
    }

    @Override
    public boolean supportsAggressiveRelease() {
        return false;
    }

    @Override
    public boolean isUnwrappableAs(Class unwrapType) {
        return unwrapType.isInstance(this);
    }

    @Override
    public <T> T unwrap(Class<T> unwrapType) {
        return (T) this;
    }
}

My CurrentTenantIdentifierResolverImpl class:

@Component
public class CurrentTenantIdentifierResolverImpl implements CurrentTenantIdentifierResolver<String> {
    @Override
    public String resolveCurrentTenantIdentifier() {
        return TenantContext.getCurrentTenant();
    }

    @Override
    public boolean validateExistingCurrentSessions() {
        return true;
    }
}

My JpaConfig class:

@Configuration
@EnableJpaRepositories(
        repositoryBaseClass = CustomJpaRepositoryImpl.class,
        repositoryFactoryBeanClass = CustomJpaRepositoryFactoryBean.class,
        basePackages = "com.mahshellsoft.bims"
)
@EnableJpaAuditing(auditorAwareRef = "auditorAware")
public class JpaConfig {
}

My CustomJpaRepositoryFactoryBean class:

public class CustomJpaRepositoryFactoryBean<R extends JpaRepository<T, ID>, T, ID extends Serializable>
        extends JpaRepositoryFactoryBean<R, T, ID> {

    public CustomJpaRepositoryFactoryBean(Class<? extends R> repositoryInterface) {
        super(repositoryInterface);
    }

    @Override
    protected RepositoryFactorySupport createRepositoryFactory(EntityManager entityManager) {
        return new CustomJpaRepositoryFactory(entityManager);
    }

    private static class CustomJpaRepositoryFactory<T, ID extends Serializable> extends JpaRepositoryFactory {

        private final EntityManager entityManager;

        public CustomJpaRepositoryFactory(EntityManager entityManager) {
            super(entityManager);
            this.entityManager = entityManager;
        }

        @Override
        protected SimpleJpaRepository<?, ?> getTargetRepository(RepositoryInformation information, EntityManager entityManager) {
            return new CustomJpaRepositoryImpl<>(information.getDomainType(), entityManager);
        }

        @Override
        protected Class<?> getRepositoryBaseClass(RepositoryMetadata metadata) {
            return CustomJpaRepositoryImpl.class;
        }
    }
}

And lastly my CustomJpaRepositoryImpl class:

public class CustomJpaRepositoryImpl<T, ID extends Serializable> extends SimpleJpaRepository<T, ID> {

    @PersistenceContext
    private EntityManager entityManager;

    public CustomJpaRepositoryImpl(Class<T> domainClass, EntityManager entityManager) {
        super(domainClass, entityManager);
        this.entityManager = entityManager;
    }

    @Override
    @Transactional
    public <S extends T> S save(S entity) {
        setSchema();
        return super.save(entity);
    }

    @Override
    @Transactional
    public void delete(T entity) {
        setSchema();
        super.delete(entity);
    }

    @Override
    @Transactional
    public void deleteById(ID id) {
        setSchema();
        super.deleteById(id);
    }

    @Override
    @Transactional
    public Optional<T> findById(ID id) {
        setSchema();
        return super.findById(id);
    }

    @Override
    @Transactional
    public List<T> findAll() {
        setSchema();
        return super.findAll();
    }

    @Override
    @Transactional
    public List<T> findAll(Sort sort) {
        setSchema();
        return super.findAll(sort);
    }

    @Override
    @Transactional
    public Page<T> findAll(Pageable pageable) {
        setSchema();
        return super.findAll(pageable);
    }

    private void setSchema() {
        String tenantId = TenantContext.getCurrentTenant();
        entityManager.createNativeQuery("SET SCHEMA '" + tenantId + "'").executeUpdate();
    }
}

This workflow works with jpa query methods like find all, save etc but not findByName,findBySchema for any repository, how do solve this. If you’ve encountered this type of problem, do share a working solution, i appreciate. Thanks