Strategies for multi-tenancy

[mike-lloyd03]

I am using postgres and am trying to figure out the best strategy for setting up my DB for a multi-tenant application using Axum (or Actix).

Currently, I am using row level access controls to require all queries to be limited to tenant_id = app.current_tenant.

CREATE POLICY tenant_isolation_policy ON users
USING (tenant_id = current_setting('app.current_tenant')::INT);
ALTER TABLE users ENABLE ROW LEVEL SECURITY;

Now when I create my DB connection, I can set the options there:

    let options = PgConnectOptions::new()
        .options([("app.current_tenant", current_tenant)]);

    let pool = PgPoolOptions::new()
        .max_connections(5)
        .connect_with(options)
        .await?;

The thing I'm not real clear on is this: can I set this value on a per-request basis? I've tried something like this:

pub async fn get_all(State(pool): State<PgPool>, query: Query<TenantId>) -> impl IntoResponse {
    let options = PgConnectOptions::new().options([("app.current_tenant", query.tenant_id)]);
    pool.set_connect_options(options);

    let users = User::get_all(&pool).await.unwrap();
    Json(users)
}

But this doesn't work because the new option is only set on subsequent DB connections.

Should I just make a new DB connection with PgPoolOptions when handling the request instead of passing a pool to the application Shared State? If I go this route, will there be a performance impact?

Is there a cleaner way to do this? Or should I just add where tenant_id = $1 to my queries and be done with it?

[mike-lloyd03]

I just did a quick test and found there's about a 2x performance hit when creating a new DB pool with every request:

pub async fn db_with_tenant_id(tenant_id: usize) -> anyhow::Result<Pool<Postgres>> {
    let pool = PgPoolOptions::new()
        .max_connections(5)
        .after_connect(move |conn, _meta| {
            Box::pin(async move {
                conn.execute(format!("SET app.current_tenant = '{}';", tenant_id).as_str())
                    .await?;

                Ok(())
            })
        })
        .connect(&DB_URL)
        .await?;

    Ok(pool)
}

pub async fn get_all() -> impl IntoResponse {
    let pool = db_with_tenant_id(1).await.unwrap();

    let users = User::get_all(&pool).await.unwrap();
    Json(users)
}

pub async fn get_all_with_pool(State(pool): State<PgPool>) -> impl IntoResponse {
    let users = User::get_all(&pool).await.unwrap();
    Json(users)
}

I then curled each endpoint 1000 times:

time for i in (seq 1000); curl 'http://localhost:8080/api/users' &> /dev/null; end

Results:
Creating a new pool on each request:

________________________________________________________
Executed in   22.99 secs    fish           external
   usr time    4.74 secs  169.99 millis    4.57 secs
   sys time    5.25 secs  547.75 millis    4.70 secs

Using a shared pool:

________________________________________________________
Executed in   11.36 secs    fish           external
   usr time    4.88 secs  138.86 millis    4.74 secs
   sys time    4.83 secs  521.27 millis    4.31 secs

[scrogson]

I believe what you want to do is run this in a transaction:

BEGIN;
SELECT set_config('app.current_tenant', $1, true);
-- the rest of your query
COMMIT;

[galah92]

I'm also interested in this. In postgres.js there's a way to reserve() specific connection from the pool, which means we reserve connection, set app.current_tenant for it, and keep it until we're done with it.
Does a similar mechanism make sense here?

[SorenHolstHansen]

What we do is we create a pg pool once like normal and set it as an extension in axum, so in your main.rs you do

let pool = PgPoolOptions::new()
        .max_connections(5)
        .connect("postgres://postgres:password@localhost/restricted_by_rls").await?;

let router = Router::new()
    // ...
    .layer(Extension(pool))

And then we have made a custom extractor for an authenticated connection like the following (You will have to define your own error stuff here)

pub struct AuthedConnection(pub PoolConnection<Postgres>);

#[async_trait]
impl<S> FromRequestParts<S> for AuthedConnection {
    type Rejection = SomeError;

    async fn from_request_parts(parts: &mut Parts, state: &S) -> Result<Self, Self::Rejection> {
        // First, get the tenant_id from the request
        // Probably something like
        // let tenant = parts.extract::<Tenant>(state).await?


        // Then get the db from the normal db extension
        let Extension(pool) = parts
            .extract::<Extension<PgPool>>()
            .await?;
        // Grab a connection
        let conn = pool.acquire().await?;
        sqlx::query("SELECT set_config('app.current_tenant', $1, false);")
            .bind(tenant)
            .execute(conn)
            .await?;
        Ok(AuthedConnection(conn))
    }
}

And then in your handlers you simply do

async fn my_handler(
    AuthedConnection(mut conn): AuthedConnection,
) -> Result<Json<MyModel>, ...> {
    let a = sqlx::query!("...").fetch_one(&mut conn).await?;

    Ok(Json(a))
}

Now the last step here is to make sqlx compile time queries happy. If you have set the DATABASE_URL to the database that is restricted by RLS, you will get errors. You have to set the DATABASE_URL to an admin url, and then the url you use to create the pool should then be a db that is restricted.

[scrogson]

I believe this sets this value for the life of the connection which could lead to some pretty gnarly security issues if not overridden in each request.

For instance, if your app is just checking out a connection from the pool and doesn’t run through this extractor, then it could be configured with the previous tenant it was configured with.

The way I’ve been thinking about this is using a transaction along with setting the third argument of current_setting to true so that it isn’t persistent outside of the current transaction.

Thoughts?

[SorenHolstHansen]

Yeah, seems good, but you can't really commit the transaction from the extractor, and sqlx calls rollback on drop. But I guess you can just remember to commit the transaction in the handler everytime

[SorenHolstHansen]

You could also just make sure to remove the config for the connection after it has been used by setting the after_release option on the pool.