pgxpool and multi-tenancy BeforeAcquire/AfterAcquire

Question

I'm trying to use postgres user variables to handle multi-tenancy through RLS. However, for performance I want my webapp to use the pg connection pool. I came across this thread: https://github.com/jackc/pgx/issues/288#issuecomment-901975396 and decided to try it.

I had to change a few functions (changes over time) but my client ends up looking like this:

package db

import (
    "context"
    "fmt"
    "github.com/jackc/pgx/v5"
    "github.com/jackc/pgx/v5/pgxpool"
    "log"
    "os"
    "time"
)

func NewPool() (pool *pgxpool.Pool, err error) {
    dsn := os.Getenv("POSTGRESQL_URL")
    config, err := pgxpool.ParseConfig(dsn)

    if err != nil {
        return pool, err
    }

    config.BeforeAcquire = func(ctx context.Context, conn *pgx.Conn) bool {
        // set the member id into this connection's setting
        memberId := ctx.Value("member_id").(string)
        _, err := conn.Exec(ctx, "select sp_set_member($1)", memberId)

        if err != nil {
            log.Fatal(err)
            return false
        } else {
            fmt.Println("Set session to memberId: " + memberId)
        }

        return true
    }

    config.AfterRelease = func(conn *pgx.Conn) bool {
        // set the setting to be empty before this connection is released to pool
        _, err := conn.Exec(context.Background(), "select sp_set_member($1)", "")

        if err != nil {
            log.Fatal(err)
            return false
        } else {
            fmt.Println("Cleared the member id")
        }

        return true
    }

    config.MaxConns = int32(20)
    config.MaxConnLifetime = time.Minute
    config.MaxConnIdleTime = time.Minute

    pool, err = pgxpool.NewWithConfig(context.Background(), config)
    return pool, err
}

Using go-chi, my route looks like this:

r.Get("/test-sql", func(w http.ResponseWriter, r *http.Request) {

            pool, err := db.NewPool()

            if err != nil {
                fmt.Println(err)
                panic(err)
            }

            rows, err := pool.Query(r.Context(), "SELECT uuid, name, owner_uuid FROM businesses")
            if err != nil {
                log.Fatal(err)
            }

            defer rows.Close()

            var rowSlice []Row
            for rows.Next() {
                var r Row

                err := rows.Scan(&r.UUID, &r.Name, &r.OwnerUUID)
                if err != nil {
                    log.Fatal(err)
                }
                rowSlice = append(rowSlice, r)
            }
            if err := rows.Err(); err != nil {
                log.Fatal(err)
            }

            fmt.Println(rowSlice)

            render.JSON(w, r, "ok sql")
        })

Watching the logs, I was expecting to see this:

app            | c7bc6655-587b-46d9-b185-2c9bd6b385fe user found
app            | Set session to memberId: c7bc6655-587b-46d9-b185-2c9bd6b385fe
app            | [{542c80c4-c067-47bb-8f25-db25dd7a184a TEST COMPANY 86c94640-dccb-4f2b-9ca4-123d38d49dca} {ce143106-fab3-4a69-8900-71a6ab2d02b3 COMPANY 2 c7bc6655-587b-46d9-b185-2c9bd6b385fe}]
app            | Cleared the member id
app            | 2022/11/26 10:58:46 [613d7c5a7d40/IiNksvJHjv-000001] "GET http://localhost:8000/api/core/test-sql HTTP/1.1" from 172.22.0.1:63464 - 200 9B in 83.631084ms

but instead I get this

app            | c7bc6655-587b-46d9-b185-2c9bd6b385fe user found
app            | Set session to memberId: c7bc6655-587b-46d9-b185-2c9bd6b385fe
app            | Cleared the member id
app            | [{542c80c4-c067-47bb-8f25-db25dd7a184a TEST COMPANY 86c94640-dccb-4f2b-9ca4-123d38d49dca} {ce143106-fab3-4a69-8900-71a6ab2d02b3 COMPANY 2 c7bc6655-587b-46d9-b185-2c9bd6b385fe}]
app            | 2022/11/26 10:58:46 [613d7c5a7d40/IiNksvJHjv-000001] "GET http://localhost:8000/api/core/test-sql HTTP/1.1" from 172.22.0.1:63464 - 200 9B in 83.631084ms

So it looks like the variable is being set and then immediately cleared...

What have I not understood here? Any help would be awesome.

Answer 1

The Pool needed to be created OUTSIDE of the route. So before the request context was created.