If you like WireQuery, please ⭐️ it on GitHub

Loading...

Blog Posts

Spring Data JDBC and Multi-Tenancy with PostgreSQL

When I was developing WireQuery, one of its requirements was to have support for multi-tenancy.

When implementing multi-tenancy, there are a number of trade-offs one needs to make, but I believe that the final implementation is exactly in the the sweet spot and one worth sharing.

But first of all, I want to acknowledge Björn Beskow's blog post on "Dynamic Multi Tenancy with Spring Boot, Hibernate and Liquibase". This blog post and the way multi-tenancy was implemented in WireQuery has been heavily inspired by this article.

What is Multi-Tenancy?

Imagine having a single application serving multiple clients, or "tenants," but keeping their data isolated and secure. That's multi-tenancy in a nutshell. It's like having a lively neighborhood where everyone has their own space, but they share the same block.

There are different ways to achieve multi-tenancy. For example:

  • Separate instance and database/schema for each tenant

    Having a separate set of instances and database/schema for each tenant has the benefit of high isolation and is easily horizontally scalable. However, the downside of this approach is the cost associated with each tenant: each instance requires a separate instance of the program to be spinned up. Maintaining the database is much more difficult, because you need to query or apply changes to each database separately.

  • Shared instance and separate schema or database for each tenant

    Having a shared instance with separate schemas or databases has the benefit of lower instance costs, while maintaining a high level of isolation, but still requires the burden of maintaining each database or schema separately.

  • Shared Databases and Discriminator Columns

    Using a shared database with discriminator columns (e.g. tenant_id) requires drastically less resources and is much easier to maintain. The downside is that it's harder to scale horizontally and that you may need to deal with "noisy neighbours" (i.e. if one tenant requires a lot of resources, this may affect other tenants as well). The benefit, however, is that you require drastically less resources (you may not need more resources for one tenant than for 100) and that the database is much easier to maintain: you only need a single database.

Having a shared Databases with discriminator columns approach is the most resource efficient. Furthermore, even if this approach is implemented, you can still combine this approach with the other approaches when needed. The only problem with this approach, is that we don't want to litter our code with queries like WHERE tenant_id = ..., as that would ripple through our entire application and introduce the risk of leaky entities, i.e. a hefty security risk.

However, if we're able to solve this issue on a more fundamental level, we get the best of both worlds: a resource efficient - yet secure solution.

Luckily, Postgres has a feature called "row level security" which can help us achieve this goal.

But how does one implement this strategy without compromising security using Spring Data JDBC (not to be confused with Spring Data JPA)? And what does it have to do with "row level security"? Let's find out!

Creating a new Spring Data JDBC Project

To illustrate how to implement multi-tenancy using Spring Data JDBC, let's build a simple customer management system. This system will have one entity, Customer, with two fields: name and address. Furthermore, this system will have two endpoints: GET /customers for retrieving customers and POST /customers for creating a new customer.

Let's start by heading over to Spring Initializr and creating a project using "Spring Web", "Flyway Migration", "Spring Data JDBC" and "PostgreSQL Driver SQL".

Next, let's create an entity, a repository and a controller:

Entity, Repository and Controller

The Customer entity:

// Customer.java

package com.example.multitenant.customer;

import org.springframework.data.annotation.Id;
import org.springframework.data.relational.core.mapping.Table;

@Table
public class Customer {
    @Id
    private Integer id;

    private String name;
    private String address;

    // ... Getters and setters
}

The Repository:

// CustomerRepository.java

package com.example.multitenant.customer;

import org.springframework.data.repository.CrudRepository;
import org.springframework.stereotype.Repository;

@Repository
public interface CustomerRepository extends CrudRepository<Customer, Integer> {
}

The Controller:

package com.example.multitenant.customer;

import org.springframework.web.bind.annotation.*;

@RestController
@RequestMapping("/customers")
public class CustomerController {
    private final CustomerRepository customerRepository;

    public CustomerController(CustomerRepository customerRepository) {
        this.customerRepository = customerRepository;
    }

    @GetMapping
    public Iterable<Customer> getCustomers() {
        return customerRepository.findAll();
    }

    @PostMapping
    public Customer createCustomer(@RequestBody Customer customer) {
        return customerRepository.save(customer);
    }
}

Setting up the Database

Next up is a docker-compose.yml file so that we can have a development Postgres database on port 5432:

version: '3'

services:
  db:
    image: postgres
    ports:
      - '5432:5432'
    environment:
      POSTGRES_PASSWORD: 'postgres'

In a separate terminal, run: docker-compose up in order to start a development database.

We need to have a separate Postgres user with support for "Row-Level Security". Wait for Postgres to be fully loaded. Then, connect to the database using postgres / postgres as your password. Then, execute the following SQL queries:

CREATE USER "mt_user" WITH PASSWORD 'mt_user';
GRANT SELECT, INSERT, UPDATE, DELETE, REFERENCES ON ALL TABLES IN SCHEMA public TO "mt_user";
GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA public TO "mt_user";
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT, INSERT, UPDATE, DELETE, REFERENCES ON TABLES TO "mt_user";
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT USAGE, SELECT ON SEQUENCES TO "mt_user";

Let's update the application.properties to ensure that our project can connect to the Postgres database:

spring.flyway.url=jdbc:postgresql://localhost:5432/postgres
spring.flyway.user=postgres
spring.flyway.password=postgres

db.url=jdbc:postgresql://localhost:5432/postgres
db.user=mt_user
db.password=mt_user

Database Migrations

Flyway provides a simple way to perform database migrations on startup. In resources/db/migration, add a file called: V1__initial-setup.sql, and add the following content:

CREATE TABLE customer
(
    id        BIGSERIAL PRIMARY KEY,
    tenant_id INT          NOT NULL DEFAULT current_setting('app.tenant_id')::INT,
    name      VARCHAR(255) NOT NULL,
    address   TEXT         NOT NULL
);

CREATE POLICY customer_tenant_isolation_policy ON customer
    USING (tenant_id = current_setting('app.tenant_id')::INT);

ALTER TABLE customer
    ENABLE ROW LEVEL SECURITY;

This is where the core of our solution lies: Postgres has a mechanism called Row Level Security which will only allow us to interact with certain data if we meet a certain condition. In this case, this condition is that the tenant_id of the customer table needs to be equal to the app.tenant_id in the current_setting (which is bound to the current database session). This value is set by the TenantAwareDataSource, which we will describe next.

Tenant Aware Configuration

The final step is to make sure our project understands different tenants. In this project, we will differentiate between tenants using the Tenant-Id header.

First, let's create a TenantAwareDataSource, which is capable to instruct Postgres to use the relevant tenant id.

The idea here is to make sure that SET app.tenant_id = '...' is called whenever a connection is obtained, whereas this value is reset whenever the connection is closed. Since Connection is only an interface and we want to support multi-tenancy for every kind of Connection, we need to create a so-called proxy.

This proxy makes sure that RESET app.tenant_id is called whenever a connection is closed.

// TenantAwareDataSource.java

package com.example.multitenant.config;

import org.jetbrains.annotations.NotNull;
import org.springframework.beans.factory.support.ScopeNotActiveException;
import org.springframework.jdbc.datasource.ConnectionProxy;
import org.springframework.jdbc.datasource.DelegatingDataSource;

import javax.sql.DataSource;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.sql.Connection;
import java.sql.SQLException;
import java.util.Objects;

public class TenantAwareDataSource extends DelegatingDataSource {
    private final TenantRequestContext tenantRequestContext;

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

    @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 {
        Integer tenantId;
        try {
            tenantId = tenantRequestContext.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);
            };
        }
    }
}

Testing Our solution

The implementation is now finished. Time to test it! Start the application and run the following command:

curl localhost:8080/customers \
    -X POST \
    -H 'Accepts: application/json' \
    -H 'Content-Type: application/json' \
    --data '{"name": "Wouter", "address": "Unknown"}' \
    -H 'Tenant-Id: 1'

Which should return:

{"id":1,"name":"Wouter","address":"Unknown"}

Now, run the same command but with Tenant-Id: 2:

curl localhost:8080/customers \
    -X POST \
    -H 'Accepts: application/json' \
    -H 'Content-Type: application/json' \
    --data '{"name": "Wouter", "address": "Unknown"}' \
    -H 'Tenant-Id: 2'

Which should return:

{"id":2,"name":"Wouter","address":"Unknown"}

Finally, let's try to fetch the customers with Tenant-Id: 1:

curl localhost:8080/customers -H 'Tenant-Id: 1'

And if all went well, we only see the first customer we created:

{"id":1,"name":"Wouter","address":"Unknown"}

Yet if we run the same command with Tenant-Id: 2, we should only see the second customer we created:

{"id":2,"name":"Wouter","address":"Unknown"}

And so, that's how you build a multi-tenant application using Spring Data JDBC and Postgres!

Conclusion

In this article, we've used the same approach of how multi-tenancy was built into WireQuery to build a simple multi-tenant customer management system.

The code for this blog post can be found at: https://github.com/wnederhof/multitenant-spring-data-jdbc.

Loading...

Back to Blog Posts