spec-kit/internal/store/postgres.go - master - seypher/benchmark

spec-kit/internal/store/postgres.go
88 строк3 KB
dorshanskiy
first_commit
19 май 2026, 18:09
19 май 2026, 18:095d12ee6
package store

import (
	"context"
	"errors"
	"fmt"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"
	"github.com/jackc/pgx/v5/pgxpool"
	"urlshortener/internal/model"
)

var (
	ErrNotFound     = errors.New("short code not found")
	ErrCodeConflict = errors.New("code conflict: try again")
)

// Store is the persistence contract for short links.
type Store interface {
	Save(ctx context.Context, code, originalURL string) (sl model.ShortLink, created bool, err error)
	FindByCode(ctx context.Context, code string) (model.ShortLink, error)
}

// PostgresStore implements Store using a pgxpool connection pool.
type PostgresStore struct {
	pool *pgxpool.Pool
}

func New(pool *pgxpool.Pool) *PostgresStore {
	return &PostgresStore{pool: pool}
}

// Save inserts a new short link or returns the existing one for the same URL.
// created is true when a new row was inserted, false when the URL already existed.
func (s *PostgresStore) Save(ctx context.Context, code, originalURL string) (model.ShortLink, bool, error) {
	var sl model.ShortLink

	// INSERT ... ON CONFLICT (original_url) DO NOTHING returns the new row only on insert.
	err := s.pool.QueryRow(ctx, `
		INSERT INTO short_links (code, original_url)
		VALUES ($1, $2)
		ON CONFLICT (original_url) DO NOTHING
		RETURNING id, code, original_url, created_at
	`, code, originalURL).Scan(&sl.ID, &sl.Code, &sl.OriginalURL, &sl.CreatedAt)

	if err == nil {
		return sl, true, nil
	}

	if errors.Is(err, pgx.ErrNoRows) {
		// URL already exists — fetch the existing record.
		err = s.pool.QueryRow(ctx, `
			SELECT id, code, original_url, created_at
			FROM short_links
			WHERE original_url = $1
		`, originalURL).Scan(&sl.ID, &sl.Code, &sl.OriginalURL, &sl.CreatedAt)
		if err != nil {
			return model.ShortLink{}, false, fmt.Errorf("store.Save lookup: %w", err)
		}
		return sl, false, nil
	}

	// Unique violation on the code column means the generated code collided.
	var pgErr *pgconn.PgError
	if errors.As(err, &pgErr) && pgErr.Code == "23505" && pgErr.ConstraintName == "uq_short_links_code" {
		return model.ShortLink{}, false, ErrCodeConflict
	}

	return model.ShortLink{}, false, fmt.Errorf("store.Save: %w", err)
}

// FindByCode returns the short link for the given code.
func (s *PostgresStore) FindByCode(ctx context.Context, code string) (model.ShortLink, error) {
	var sl model.ShortLink
	err := s.pool.QueryRow(ctx, `
		SELECT id, code, original_url, created_at
		FROM short_links
		WHERE code = $1
	`, code).Scan(&sl.ID, &sl.Code, &sl.OriginalURL, &sl.CreatedAt)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return model.ShortLink{}, ErrNotFound
		}
		return model.ShortLink{}, fmt.Errorf("store.FindByCode: %w", err)
	}
	return sl, nil
}