[Question] Any way to use functions from pgx.Tx and *pgxpool.Pool without if-else check?

[bentcoder]

Hi,

As you can see below, given *pgxpool.Pool and pgx.Tx are two different types I have if-else to decide what to use to run a query. The problem here is that, I have so many methods (for Query and Exec too) do the same thing (if-else) and this is not maintainable in long run. What would be the cleanest way to refactor the code to get rid of if-else repetition in all methods?

Thank you

import (
	"context"

	"github.com/jackc/pgx/v5/pgxpool"
)

type Storage struct {
	dtb *pgxpool.Pool
}

func (s Storage) CountUsers(ctx context.Context, id string, txn *pgx.Tx) (int, error) {
	qry := `SELECT COUNT(*) FROM users WHERE id = $1`

	var (
		err   error
		count int
	)

	if txn != nil {
		err = txn.QueryRow(ctx, qry, id).Scan(&count)
	} else {
		err = s.dtb.QueryRow(ctx, qry, id).Scan(&count)
	}

	if err != nil {
		return 0, err
	}

	return count, nil
}

func (s Storage) GetUsers(ctx context.Context, ..., txn *pgx.Tx) (int, error) {
	// ...

	if txn != nil {
		err = txn.Query(ctx, ...)
	} else {
		err = s.dtb.QueryRow(ctx, ...)
	}

	// ...
}

func (s Storage) CreateUser(ctx context.Context, ..., txn *pgx.Tx) (int, error) {
	// ...

	if txn != nil {
		err = txn.Exec(ctx, ...)
	} else {
		err = s.dtb.Exec(ctx, ...)
	}

	// ...
}

func (s Storage) ...

[Jack Christensen]

What I do is have my functions take a db value that is an interface that pools, connections, and transactions all satisfy. Then the functions are implemented in terms of that interface.

I would suggest that structure. However, if you really want to keep the existing method signatures then I would still define the common interface and create a method on storage that took a *pgx.Tx and returned that interface. If the tx was present it would return it and otherwise it would return the pool member of storage. This method could be called at the top of each method that takes a tx and then the rest of the method could be written in terms of that interface.

[bentcoder]

Thank you @jackc for the response. I wouldn't pass *pgxpool.Pool to my methods as dependency/argument because my Storage type already has it at application boot. Hence you first sentence wouldn't apply to me for the current design I have.

Although I have a feeling it might be promising, I am not perfectly sure if I understand the solution you mentioned in the second sentence. If that won't be too much for you, do you mind showing an example please?

[Jack Christensen]

Something like this:

type DB interface {
	Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error)
	Query(ctx context.Context, sql string, optionsAndArgs ...interface{}) (pgx.Rows, error)
	QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
}

func (s Storage) GetDBOrTx(txn *pgx.Tx) DB {
	if txn != nil {
		return txn
	} else {
		return s.dtb
	}
}

func (s Storage) CreateUser(ctx context.Context, ..., txn *pgx.Tx) (int, error) {
	db := s.GetDBOrTx(txn)
	err = db.Exec(ctx, ...)
	// ...
}

[bentcoder]

Thank you @jackc, much appreciated. One question. Why do you pass pointer txn *pgx.Tx which is an interface? Maybe a typo? If not, that will produce error below.

cannot use txn (variable of type *pgx.Tx) as DB value in return statement: *pgx.Tx does not implement DB
(type *pgx.Tx is pointer to interface, not interface)compiler[InvalidIfaceAssign]
(https://pkg.go.dev/golang.org/x/tools/internal/typesinternal#InvalidIfaceAssign)
var txn *pgx.Tx

[Jack Christensen]

Why do you pass pointer txn *pgx.Tx which is an interface? Maybe a typo?

Right. That was a typo. Sorry.

[WhiskeyJack96]

@bentcoder another option is to have storage hold on to an interface reference and have a "Tranasactional" method. I typically prefer the strategy @jackc already suggested (interface for the methods) but this is also viable if you want to keep youre method signatures "db agnostic"

type DB interface {
	Exec(ctx context.Context, sql string, arguments ...any) (pgconn.CommandTag, error)
	Query(ctx context.Context, sql string, optionsAndArgs ...interface{}) (pgx.Rows, error)
	QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
}

type Storage struct {
	dtb DB
}

func (s Storage) Transactionally(txn pgx.Tx) Storage {
	return Storage{dbt: txn}
}

func (s Storage) CreateUser(ctx context.Context, ...) (int, error) {
	err = s.dtb.Exec(ctx, ...)
	// ...
}

[bentcoder]

@WhiskeyJack96 That's literally what I am doing but the only difference is that I don't recreate Storage in transaction method. Instead, I just return either txn or pool itself using if.

[bentcoder]

@WhiskeyJack96 How you obtain TXN from your caller before passing to Transactionally(txn pgx.Tx)? I am hoping you are not passing *pgxpool.Pool to your caller as dependency. Saying it because I am trying to improve my code where I don't have *pgxpool.Pool in the caller and don't want to. Here is how my code looks and it feels a bit cluttered.

@jackc You are more than welcome if you want to thrown in input.

My aim here is that I want to avoid passing txn pgx.Tx to every single method.

package main

pool := postgres.NewPool()

str := storage.Storage{
	Database: pool,
}

svc := user.Service{
	Storage: str,
}

...

package user

type storer interface {
	Transaction(ctx context.Context) (pgx.Tx, error)
	DeleteUser(ctx context.Context, txn pgx.Tx, ID string) error
	DeleteAllPostsByUser(ctx context.Context, txn pgx.Tx, userID string) error
}

type User struct {
	Storage storer
}

func (u User) Delete(ctx context.Context) error {
	txn, err := u.Storage.Transaction(ctx)
	if err != nil {
		return err
	}
	defer txn.Rollback(ctx)

	if err := u.Storage.DeleteUser(ctx, txn, "xyz"); err != nil {
		return err
	}

	if err := u.Storage.DeleteAllPostsByUser(ctx, txn, "xyz"); err != nil {
		return err
	}

	return txn.Commit(ctx)
}

package storage

type querier interface {
	Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error)
	Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error)
	QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
}

type Storage struct {
	Database *pgxpool.Pool
}

func (s Storage) Transaction(ctx context.Context) (pgx.Tx, error) {
	return s.Database.Begin(ctx)
}

func (s Storage) with(txn pgx.Tx) querier {
	if txn != nil {
		return txn
	}

	return s.Database
}

func (s Storage) DeleteUser(ctx context.Context, txn pgx.Tx, ID string) error {
	s.with(txn).Exec(ctx, qry, id)
	...
}

func (s Storage) DeleteAllPostsByUser(ctx context.Context, txn pgx.Tx, userID string) error {
	s.with(txn).Exec(ctx, qry, userID)
	...
}

func (s Storage) SomeOtherMethod(ctx context.Context, txn pgx.Tx, ...) error {
	...
}