payments/internal/storage/workflow_instances.go

package storage

import (
	"context"
	"fmt"
	"time"

	"github.com/hanzo-fi/go-libs/v5/pkg/query"
	"github.com/hanzo-fi/go-libs/v5/pkg/storage/bun/paginate"
	"github.com/hanzo-fi/go-libs/v5/pkg/types/pointer"
	internalTime "github.com/hanzo-fi/go-libs/v5/pkg/types/time"
	"github.com/hanzo-fi/payments/pkg/domain/models"
	"github.com/uptrace/bun"
)

type instance struct {
	bun.BaseModel `bun:"table:workflows_instances"`

	// Mandatory fields
	ID          string             `bun:"id,pk,type:text,notnull"`
	ScheduleID  string             `bun:"schedule_id,pk,type:text,notnull"`
	ConnectorID models.ConnectorID `bun:"connector_id,pk,type:character varying,notnull"`
	CreatedAt   internalTime.Time  `bun:"created_at,type:timestamp without time zone,notnull"`
	UpdatedAt   internalTime.Time  `bun:"updated_at,type:timestamp without time zone,notnull"`

	// Optional fields with default
	// c.f. https://bun.uptrace.dev/guide/models.html#default
	Terminated bool `bun:"terminated,type:boolean,notnull,nullzero,default:false"`

	// Optional fields
	// c.f.: https://bun.uptrace.dev/guide/models.html#nulls
	TerminatedAt *internalTime.Time `bun:"terminated_at,type:timestamp without time zone,nullzero"`
	Error        *string            `bun:"error,type:text,nullzero"`
}

func (s *store) InstancesUpsert(ctx context.Context, instance models.Instance) error {
	toInsert := fromInstanceModel(instance)

	_, err := s.db.NewInsert().
		Model(&toInsert).
		On("CONFLICT (id, schedule_id, connector_id) DO NOTHING").
		Exec(ctx)

	return e("failed to insert new instance", err)
}

func (s *store) InstancesUpdate(ctx context.Context, instance models.Instance) error {
	toUpdate := fromInstanceModel(instance)

	_, err := s.db.NewUpdate().
		Model(&toUpdate).
		Set("updated_at = ?", instance.UpdatedAt).
		Set("terminated = ?", instance.Terminated).
		Set("terminated_at = ?", instance.TerminatedAt).
		Set("error = ?", instance.Error).
		WherePK().
		Exec(ctx)

	return e("failed to update instance", err)
}

func (s *store) InstancesDeleteFromConnectorID(ctx context.Context, connectorID models.ConnectorID) error {
	_, err := s.db.NewDelete().
		Model((*instance)(nil)).
		Where("connector_id = ?", connectorID).
		Exec(ctx)

	return e("failed to delete instances", err)
}

// InstancesDeleteFromConnectorIDBatch deletes a batch of instances for a given connector ID
// and returns the number of rows affected
func (s *store) InstancesDeleteFromConnectorIDBatch(ctx context.Context, connectorID models.ConnectorID, batchSize int) (int, error) {
	if batchSize <= 0 {
		return 0, fmt.Errorf("invalid batchSize %d for connectorID %s: %w", batchSize, connectorID.String(), ErrValidation)
	}

	result, err := s.db.NewDelete().
		Model((*instance)(nil)).
		Where("(id, schedule_id, connector_id) IN (SELECT id, schedule_id, connector_id FROM workflows_instances WHERE connector_id = ? LIMIT ?)", connectorID, batchSize).
		Exec(ctx)
	if err != nil {
		return 0, e("failed to delete instances batch", err)
	}

	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return 0, e("failed to get rows affected", err)
	}

	return int(rowsAffected), nil
}

type InstanceQuery struct{}

type ListInstancesQuery paginate.OffsetPaginatedQuery[paginate.PaginatedQueryOptions[InstanceQuery]]

func NewListInstancesQuery(opts paginate.PaginatedQueryOptions[InstanceQuery]) ListInstancesQuery {
	return ListInstancesQuery{
		Order:    paginate.OrderAsc,
		PageSize: opts.PageSize,
		Options:  opts,
	}
}

func (s *store) instancesQueryContext(qb query.Builder) (string, []any, error) {
	return qb.Build(query.ContextFn(func(key, operator string, value any) (string, []any, error) {
		switch key {
		case "schedule_id", "connector_id":
			if operator != "$match" {
				return "", nil, fmt.Errorf("'connector_id' column can only be used with $match: %w", ErrValidation)
			}
			return fmt.Sprintf("%s = ?", key), []any{value}, nil
		default:
			return "", nil, fmt.Errorf("unknown key '%s' when building query: %w", key, ErrValidation)
		}
	}))
}

func (s *store) InstancesGet(ctx context.Context, id string, scheduleID string, connectorID models.ConnectorID) (*models.Instance, error) {
	var i instance
	err := s.db.NewSelect().
		Model(&i).
		Where("id = ?", id).
		Where("schedule_id = ?", scheduleID).
		Where("connector_id = ?", connectorID).
		Scan(ctx)
	if err != nil {
		return nil, e("failed to fetch instance", err)
	}

	return pointer.For(toInstanceModel(i)), nil
}

func (s *store) InstancesList(ctx context.Context, q ListInstancesQuery) (*paginate.Cursor[models.Instance], error) {
	var (
		where string
		args  []any
		err   error
	)
	if q.Options.QueryBuilder != nil {
		where, args, err = s.instancesQueryContext(q.Options.QueryBuilder)
		if err != nil {
			return nil, err
		}
	}

	cursor, err := paginateWithOffset[paginate.PaginatedQueryOptions[InstanceQuery], instance](s, ctx,
		(*paginate.OffsetPaginatedQuery[paginate.PaginatedQueryOptions[InstanceQuery]])(&q),
		func(query *bun.SelectQuery) *bun.SelectQuery {
			if where != "" {
				query = query.Where(where, args...)
			}

			query = query.Order("created_at DESC", "sort_id DESC")

			return query
		},
	)
	if err != nil {
		return nil, e("failed to fetch instances", err)
	}

	instances := make([]models.Instance, 0, len(cursor.Data))
	for _, i := range cursor.Data {
		instances = append(instances, toInstanceModel(i))
	}

	return &paginate.Cursor[models.Instance]{
		PageSize: cursor.PageSize,
		HasMore:  cursor.HasMore,
		Previous: cursor.Previous,
		Next:     cursor.Next,
		Data:     instances,
	}, nil
}

func fromInstanceModel(from models.Instance) instance {
	return instance{
		ID:          from.ID,
		ScheduleID:  from.ScheduleID,
		ConnectorID: from.ConnectorID,
		CreatedAt:   internalTime.New(from.CreatedAt),
		UpdatedAt:   internalTime.New(from.UpdatedAt),
		Terminated:  from.Terminated,
		TerminatedAt: func() *internalTime.Time {
			if from.TerminatedAt == nil {
				return nil
			}
			return pointer.For(internalTime.New(*from.TerminatedAt))
		}(),
		Error: from.Error,
	}
}

func toInstanceModel(from instance) models.Instance {
	return models.Instance{
		ID:          from.ID,
		ScheduleID:  from.ScheduleID,
		ConnectorID: from.ConnectorID,
		CreatedAt:   from.CreatedAt.Time,
		UpdatedAt:   from.UpdatedAt.Time,
		Terminated:  from.Terminated,
		TerminatedAt: func() *time.Time {
			if from.TerminatedAt == nil {
				return nil
			}
			return pointer.For(from.TerminatedAt.Time)
		}(),
		Error: from.Error,
	}
}

// Returns the latest workflow instance for schedules which in their last N executions returned exclusively errors,
// where N is determined by the threshold parameter.
func (s *store) InstancesListSchedulesAboveErrorThreshold(ctx context.Context, connectorID models.ConnectorID, threshold int, q ListInstancesQuery) (*paginate.Cursor[models.Instance], error) {
	cursor, err := paginateWithOffset[paginate.PaginatedQueryOptions[InstanceQuery], instance](s, ctx,
		(*paginate.OffsetPaginatedQuery[paginate.PaginatedQueryOptions[InstanceQuery]])(&q),
		func(query *bun.SelectQuery) *bun.SelectQuery {
			return query.
				DistinctOn("schedule_id").
				Where("connector_id = ?", connectorID).
				Where("terminated = TRUE").
				Where(`schedule_id IN (
					SELECT schedule_id
					FROM (
						SELECT
							schedule_id,
							ROW_NUMBER() OVER (PARTITION BY schedule_id ORDER BY created_at DESC) AS rn,
							error
						FROM workflows_instances
						WHERE connector_id = ?
					) ranked
					WHERE rn <= ?
					GROUP BY schedule_id
					HAVING COUNT(*) = ? AND COUNT(*) FILTER (WHERE error IS NULL) = 0
				)`, connectorID, threshold, threshold).
				Order("schedule_id ASC", "created_at DESC")
		},
	)
	if err != nil {
		return nil, e("failed to fetch instance errors", err)
	}

	instances := make([]models.Instance, 0, len(cursor.Data))
	for _, i := range cursor.Data {
		instances = append(instances, toInstanceModel(i))
	}

	return &paginate.Cursor[models.Instance]{
		PageSize: cursor.PageSize,
		HasMore:  cursor.HasMore,
		Previous: cursor.Previous,
		Next:     cursor.Next,
		Data:     instances,
	}, nil
}