pgxpool 和多租户 BeforeAcquire/AfterAcquire

Question

我正在尝试使用 postgres 用户变量通过 RLS 处理多租户。但是，为了性能，我希望我的 webapp 使用 pg 连接池。我遇到了这个线程：https://github.com/jackc/pgx/issues/288#issuecomment-901975396 并决定尝试一下。

我不得不更改一些功能（随时间变化），但我的客户最终看起来像这样：

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
}

使用 go-chi，我的路线如下所示：

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")
        })

看着日志，我期待看到这个：

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

但我得到了这个

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

所以看起来变量正在设置然后立即清除......

我在这里有什么不明白的？任何帮助都是极好的。

Answer 1

池需要在路线之外创建。所以在创建请求上下文之前。