ms-tester/internal/problems/repository/pg_repository.go

package repository

import (
	"context"
	"errors"
	"git.sch9.ru/new_gate/models"
	"git.sch9.ru/new_gate/ms-tester/pkg/utils"
	"github.com/jackc/pgerrcode"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jmoiron/sqlx"
	"go.uber.org/zap"
)

type ProblemRepository struct {
	db     *sqlx.DB
	logger *zap.Logger
}

func NewProblemRepository(db *sqlx.DB, logger *zap.Logger) *ProblemRepository {
	return &ProblemRepository{
		db:     db,
		logger: logger,
	}
}

func (storage *ProblemRepository) CreateProblem(ctx context.Context, problem *models.Problem, testGroupData []models.TestGroupData) (int32, error) {
	tx, err := storage.db.Beginx()
	if err != nil {
		return 0, handlePgErr(err)
	}
	query := tx.Rebind(`
INSERT INTO problems
    (name,description,time_limit,memory_limit)
VALUES (?, ?, ?, ?)
RETURNING id
`)
	rows, err := tx.QueryxContext(
		ctx,
		query,
		problem.Name,
		problem.Description,
		problem.TimeLimit,
		problem.MemoryLimit,
	)
	if err != nil {
		return 0, handlePgErr(errors.Join(err, tx.Rollback()))
	}
	for _, tgd := range testGroupData {
		query := tx.Rebind(`
    INSERT INTO testgroups
        (problem_id,testing_strategy)
    VALUES ((select last_value from problems_id_seq),?)
    RETURNING id
    `)
		rows, err = tx.QueryxContext(ctx, query, tgd.Ts)
		if err != nil {
			return 0, handlePgErr(errors.Join(err, tx.Rollback()))
		}
		var i int32 = 0
		for ; i < tgd.TestAmount; i++ {
			query := tx.Rebind(`
        INSERT INTO tests
            (testgroup_id)
        VALUES ((select last_value from testgroups_id_seq))
        RETURNING id
        `)
			rows, err = tx.QueryxContext(ctx, query, tgd.Ts)
			if err != nil {
				return 0, handlePgErr(errors.Join(err, tx.Rollback()))
			}
		}
	}
	err = tx.Commit()
	//add test saving

	defer rows.Close()
	var id int32
	err = rows.StructScan(&id)
	if err != nil {
		return 0, handlePgErr(err)
	}

	return id, nil
}

func (storage *ProblemRepository) ReadProblemById(ctx context.Context, id int32) (*models.Problem, error) {
	var problem models.Problem
	query := storage.db.Rebind("SELECT * from problems WHERE id=? LIMIT 1")
	err := storage.db.GetContext(ctx, &problem, query, id)
	if err != nil {
		return nil, handlePgErr(err)
	}
	return &problem, nil
}

func (storage *ProblemRepository) UpdateProblem(ctx context.Context, problem *models.Problem) error {
	query := storage.db.Rebind("UPDATE problems SET name=?,description=?,time_limit=?,memory_limit=? WHERE id=?")
	_, err := storage.db.ExecContext(ctx, query, problem.Name, problem.Description, problem.TimeLimit, problem.MemoryLimit, problem.Id)
	if err != nil {
		return handlePgErr(err)
	}
	return nil
}

func (storage *ProblemRepository) DeleteProblem(ctx context.Context, id int32) error {
	query := storage.db.Rebind("DELETE FROM problems WHERE id=?")
	_, err := storage.db.ExecContext(ctx, query, id)
	if err != nil {
		return handlePgErr(err)
	}

	return nil
}

func handlePgErr(err error) error {
	var pgErr *pgconn.PgError
	if !errors.As(err, &pgErr) {
		return utils.StorageError(err, utils.ErrUnknown, "unexpected error from postgres")
	}
	if pgerrcode.IsIntegrityConstraintViolation(pgErr.Code) {
		// TODO: probably should specify which constraint
		return utils.StorageError(err, utils.ErrConflict, pgErr.Message)
	}
	if pgerrcode.IsNoData(pgErr.Code) {
		return utils.StorageError(err, utils.ErrNotFound, pgErr.Message)
	}
	return utils.StorageError(err, utils.ErrUnimplemented, "unimplemented error")
}