postgres.go

// Package postgres is the implementation of the postgres data store.
package postgres

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"time"

	"github.com/italolelis/outboxer"
	"github.com/italolelis/outboxer/lock"
)

const (
	// DefaultEventStoreTable is the default table name.
	DefaultEventStoreTable = "event_store"
)

var (
	// ErrLocked is used when we can't acquire an explicit lock.
	ErrLocked = errors.New("can't acquire lock")

	// ErrNoDatabaseName is used when the database name is blank.
	ErrNoDatabaseName = errors.New("no database name")

	// ErrNoSchema is used when the schema name is blank.
	ErrNoSchema = errors.New("no schema")
)

// Postgres is the implementation of the data store.
type Postgres struct {
	conn            *sql.Conn
	DatabaseName    string
	SchemaName      string
	EventStoreTable string
	isLocked        bool
}

// WithInstance creates a postgres data store with an existing db connection.
func WithInstance(ctx context.Context, db *sql.DB) (*Postgres, error) {
	conn, err := db.Conn(ctx)
	if err != nil {
		return nil, err
	}

	p := Postgres{conn: conn}

	if err := conn.QueryRowContext(ctx, `SELECT CURRENT_DATABASE()`).Scan(&p.DatabaseName); err != nil {
		return nil, err
	}

	if p.DatabaseName == "" {
		return nil, ErrNoDatabaseName
	}

	if err := conn.QueryRowContext(ctx, `SELECT CURRENT_SCHEMA()`).Scan(&p.SchemaName); err != nil {
		return nil, err
	}

	if p.SchemaName == "" {
		return nil, ErrNoSchema
	}

	if p.EventStoreTable == "" {
		p.EventStoreTable = DefaultEventStoreTable
	}

	if err := p.ensureTable(ctx); err != nil {
		return nil, err
	}

	return &p, nil
}

// Close closes the db connection.
func (p *Postgres) Close() error {
	if err := p.conn.Close(); err != nil {
		return fmt.Errorf("failed to close connection: %w", err)
	}

	return nil
}

// GetEvents retrieves all the relevant events.
func (p *Postgres) GetEvents(ctx context.Context, batchSize int32) ([]*outboxer.OutboxMessage, error) {
	events := make([]*outboxer.OutboxMessage, 0, batchSize)

	// nolint
	rows, err := p.conn.QueryContext(ctx, fmt.Sprintf("SELECT * FROM %s WHERE dispatched = false LIMIT %d", p.EventStoreTable, batchSize))
	if err != nil {
		return events, fmt.Errorf("failed to get messages from the store: %w", err)
	}

	for rows.Next() {
		var e outboxer.OutboxMessage

		err = rows.Scan(&e.ID, &e.Dispatched, &e.DispatchedAt, &e.Payload, &e.Options, &e.Headers)
		if err != nil {
			return events, fmt.Errorf("failed to scan message: %w", err)
		}

		events = append(events, &e)
	}

	return events, nil
}

// Add adds the message to the data store.
func (p *Postgres) Add(ctx context.Context, evt *outboxer.OutboxMessage) error {
	tx, err := p.conn.BeginTx(ctx, &sql.TxOptions{})
	if err != nil {
		return fmt.Errorf("transaction start failed: %w", err)
	}

	// nolint
	query := fmt.Sprintf(`INSERT INTO %s (payload, options, headers) VALUES ($1, $2, $3)`, p.EventStoreTable)
	if _, err := tx.ExecContext(ctx, query, evt.Payload, evt.Options, evt.Headers); err != nil {
		if err := tx.Rollback(); err != nil {
			return err
		}

		return fmt.Errorf("failed to insert message into the data store: %w", err)
	}

	if err := tx.Commit(); err != nil {
		return fmt.Errorf("transaction commit failed: %w", err)
	}

	return nil
}

// AddWithinTx creates a transaction and then tries to execute anything within it.
func (p *Postgres) AddWithinTx(ctx context.Context, evt *outboxer.OutboxMessage, fn func(outboxer.ExecerContext) error) error {
	tx, err := p.conn.BeginTx(ctx, &sql.TxOptions{})
	if err != nil {
		return fmt.Errorf("transaction start failed: %w", err)
	}

	if err := fn(tx); err != nil {
		return err
	}

	// nolint
	query := fmt.Sprintf(`INSERT INTO %s (payload, options, headers) VALUES ($1, $2, $3)`, p.EventStoreTable)
	if _, err := tx.ExecContext(ctx, query, evt.Payload, evt.Options, evt.Headers); err != nil {
		if err := tx.Rollback(); err != nil {
			return err
		}

		return fmt.Errorf("failed to insert message into the data store: %w", err)
	}

	if err := tx.Commit(); err != nil {
		return fmt.Errorf("transaction commit failed: %w", err)
	}

	return nil
}

// SetAsDispatched sets one message as dispatched.
func (p *Postgres) SetAsDispatched(ctx context.Context, id int64) error {
	query := fmt.Sprintf(`
update %s
set
    dispatched = true,
    dispatched_at = now(),
    options = '{}',
	headers = '{}'
where id = $1;
`, p.EventStoreTable)
	if _, err := p.conn.ExecContext(ctx, query, id); err != nil {
		return fmt.Errorf("failed to set message as dispatched: %w", err)
	}

	return nil
}

// Remove removes old messages from the data store.
func (p *Postgres) Remove(ctx context.Context, dispatchedBefore time.Time, batchSize int32) error {
	tx, err := p.conn.BeginTx(ctx, &sql.TxOptions{})
	if err != nil {
		return fmt.Errorf("transaction start failed: %w", err)
	}

	q := `
DELETE FROM %[1]s
WHERE ctid IN
(
    select ctid
    from %[1]s
    where
        "dispatched" = true and
        "dispatched_at" < $1
    limit %d
)
`

	query := fmt.Sprintf(q, p.EventStoreTable, batchSize)
	if _, err := tx.ExecContext(ctx, query, dispatchedBefore); err != nil {
		if err := tx.Rollback(); err != nil {
			return err
		}

		return fmt.Errorf("failed to remove messages from the data store: %w", err)
	}

	if err := tx.Commit(); err != nil {
		return fmt.Errorf("transaction commit failed: %w", err)
	}

	return nil
}

// Lock implements explicit locking.
// https://www.postgresql.org/docs/9.6/static/explicit-locking.html#ADVISORY-LOCKS
func (p *Postgres) lock(ctx context.Context) error {
	if p.isLocked {
		return ErrLocked
	}

	aid, err := lock.Generate(p.DatabaseName, p.SchemaName)
	if err != nil {
		return err
	}

	// This will either obtain the lock immediately and return true,
	// or return false if the lock cannot be acquired immediately.
	query := `SELECT pg_advisory_lock($1)`
	if _, err := p.conn.ExecContext(ctx, query, aid); err != nil {
		return fmt.Errorf("try lock failed: %w", err)
	}

	p.isLocked = true

	return nil
}

// Unlock is the implementation of the unlock for explicit locking.
func (p *Postgres) unlock(ctx context.Context) error {
	if !p.isLocked {
		return nil
	}

	aid, err := lock.Generate(p.DatabaseName, p.SchemaName)
	if err != nil {
		return err
	}

	query := `SELECT pg_advisory_unlock($1)`
	if _, err := p.conn.ExecContext(ctx, query, aid); err != nil {
		return err
	}

	p.isLocked = false

	return nil
}

func (p *Postgres) ensureTable(ctx context.Context) (err error) {
	if err := p.lock(ctx); err != nil {
		return err
	}

	defer func() {
		if e := p.unlock(ctx); e != nil {
			if err == nil {
				err = e
			} else {
				err = fmt.Errorf("failed to unlock table: %w", err)
			}
		}
	}()

	query := fmt.Sprintf(`
CREATE TABLE IF NOT EXISTS %[1]s (
	id SERIAL not null primary key, 
	dispatched boolean not null default false, 
	dispatched_at timestamp,
	payload bytea not null,
	options jsonb,
	headers jsonb
);

CREATE INDEX IF NOT EXISTS "index_dispatchedAt" ON %[1]s using btree (dispatched_at asc nulls last);
CREATE INDEX IF NOT EXISTS "index_dispatched" ON %[1]s using btree (dispatched asc nulls last);
`, p.EventStoreTable)

	if _, err = p.conn.ExecContext(ctx, query); err != nil {
		return err
	}

	return nil
}