Hooks to execute sql queries just after getting a connection and before releasing the connection.

[ankit-wasnik-pcn]

I want to implement row level security. Is there a way to execute the following query just after getting the connection and before releasing the connection.
e.ExecContext(ctx, "SELECT set_user($1)", "some-value")

Database: Postgres
Driver: github.com/jackc/pgx/v5

[stephenafamo]

This sort of has nothing to do with Bob 😅.

You would have to either find a driver that allows these hooks, or wrap an existing driver to hook into the connection process.

[ankit-wasnik-pcn]

If anyone is using "github.com/jackc/pgx/v5", refer to the following code. The following code creates a new *sql.DB from the pool.

import (
	"context"
	"fmt"
	"log/slog"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
	"github.com/jackc/pgx/v5/stdlib"
	"github.com/stephenafamo/bob"
)

func NewBobDb(dsn string) (bob.DB, error) {
	if dsn == "" {
		return bob.DB{}, fmt.Errorf("dsn is empty")
	}

	pool, err := NewConnPool(dsn)
	if err != nil {
		return bob.DB{}, err
	}
	sqlDB := stdlib.OpenDBFromPool(pool)
	return bob.NewDB(sqlDB), nil
}

func NewConnPool(dsn string) (pool *pgxpool.Pool, err error) {
	config, err := pgxpool.ParseConfig(dsn)
	if err != nil {
		return nil, err
	}
	config.BeforeAcquire = func(ctx context.Context, conn *pgx.Conn) bool {
		// set the user id into this connection's setting
		userId := ctx.Value("userId").(string)
		_, err := conn.Exec(ctx, "select set_user($1)", userId)
		if err != nil {
			slog.Error("failed to set user id in db session", "userId", userId, "err", err)
                        return false
		}
		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 set_user($1)", "")
		if err != nil {
			slog.Error("failed to reset user id in db session", "err", err)
                        return false
		}
		return true
	}
	config.MaxConns = int32(20)
	config.MaxConnLifetime = time.Minute
	config.MaxConnIdleTime = time.Minute
	pool, err = pgxpool.NewWithConfig(context.Background(), config)
	if err != nil {
		return nil, err
	}
	
	return pool, err
}