emissary/internal/datastore/agent_repository.go

package datastore

import (
	"context"
	"time"

	"github.com/lestrrat-go/jwx/v2/jwk"
)

type AgentRepository interface {
	Create(ctx context.Context, thumbprint string, keySet jwk.Set, metadata map[string]any) (*Agent, error)

	Attach(ctx context.Context, tenantID TenantID, agentID AgentID) (*Agent, error)
	Detach(ctx context.Context, agentID AgentID) (*Agent, error)

	Get(ctx context.Context, id AgentID) (*Agent, error)
	Update(ctx context.Context, id AgentID, updates ...AgentUpdateOptionFunc) (*Agent, error)
	Query(ctx context.Context, opts ...AgentQueryOptionFunc) ([]*Agent, int, error)
	Delete(ctx context.Context, id AgentID) error

	UpdateSpec(ctx context.Context, id AgentID, name string, version string, revision int, data map[string]any) (*Spec, error)
	QuerySpecs(ctx context.Context, id AgentID) ([]*SpecHeader, error)
	GetSpec(ctx context.Context, id AgentID, name string, version string) (*Spec, error)
	DeleteSpec(ctx context.Context, id AgentID, name string, version string) error
}

type AgentQueryOptionFunc func(*AgentQueryOptions)

type AgentQueryOptions struct {
	Limit       *int
	Offset      *int
	IDs         []AgentID
	TenantIDs   []TenantID
	Thumbprints []string
	Metadata    *map[string]any
	Statuses    []AgentStatus
}

func WithAgentQueryLimit(limit int) AgentQueryOptionFunc {
	return func(opts *AgentQueryOptions) {
		opts.Limit = &limit
	}
}

func WithAgentQueryOffset(offset int) AgentQueryOptionFunc {
	return func(opts *AgentQueryOptions) {
		opts.Offset = &offset
	}
}

func WithAgentQueryMetadata(metadata map[string]any) AgentQueryOptionFunc {
	return func(opts *AgentQueryOptions) {
		opts.Metadata = &metadata
	}
}

func WithAgentQueryID(ids ...AgentID) AgentQueryOptionFunc {
	return func(opts *AgentQueryOptions) {
		opts.IDs = ids
	}
}

func WithAgentQueryTenantID(ids ...TenantID) AgentQueryOptionFunc {
	return func(opts *AgentQueryOptions) {
		opts.TenantIDs = ids
	}
}

func WithAgentQueryStatus(statuses ...AgentStatus) AgentQueryOptionFunc {
	return func(opts *AgentQueryOptions) {
		opts.Statuses = statuses
	}
}

func WithAgentQueryThumbprints(thumbprints ...string) AgentQueryOptionFunc {
	return func(opts *AgentQueryOptions) {
		opts.Thumbprints = thumbprints
	}
}

type AgentUpdateOptionFunc func(*AgentUpdateOptions)

type AgentUpdateOptions struct {
	Label       *string
	Status      *AgentStatus
	ContactedAt *time.Time
	Metadata    *map[string]any
	KeySet      *jwk.Set
	Thumbprint  *string
	TenantID    *TenantID
}

func WithAgentUpdateTenant(id TenantID) AgentUpdateOptionFunc {
	return func(opts *AgentUpdateOptions) {
		opts.TenantID = &id
	}
}

func WithAgentUpdateStatus(status AgentStatus) AgentUpdateOptionFunc {
	return func(opts *AgentUpdateOptions) {
		opts.Status = &status
	}
}

func WithAgentUpdateMetadata(metadata map[string]any) AgentUpdateOptionFunc {
	return func(opts *AgentUpdateOptions) {
		opts.Metadata = &metadata
	}
}

func WithAgentUpdateKeySet(keySet jwk.Set) AgentUpdateOptionFunc {
	return func(opts *AgentUpdateOptions) {
		opts.KeySet = &keySet
	}
}

func WithAgentUpdateThumbprint(thumbprint string) AgentUpdateOptionFunc {
	return func(opts *AgentUpdateOptions) {
		opts.Thumbprint = &thumbprint
	}
}

func WithAgentUpdateLabel(label string) AgentUpdateOptionFunc {
	return func(opts *AgentUpdateOptions) {
		opts.Label = &label
	}
}

func WithAgentUpdateContactedAt(contactedAt time.Time) AgentUpdateOptionFunc {
	return func(opts *AgentUpdateOptions) {
		opts.ContactedAt = &contactedAt
	}
}