pkg/framework/entutils/mixins.go

package entutils

import (
	"fmt"
	"time"

	"entgo.io/ent"
	"entgo.io/ent/dialect"
	"entgo.io/ent/dialect/entsql"
	"entgo.io/ent/schema/field"
	"entgo.io/ent/schema/index"
	"entgo.io/ent/schema/mixin"
	"github.com/oklog/ulid/v2"

	"github.com/openmeterio/openmeter/pkg/clock"
	"github.com/openmeterio/openmeter/pkg/convert"
	"github.com/openmeterio/openmeter/pkg/models"
)

// UniqueResourceMixin adds common fields
type UniqueResourceMixin struct {
	mixin.Schema
}

func (UniqueResourceMixin) Fields() []ent.Field {
	fields := ResourceMixin{}.Fields()
	fields = append(fields, KeyMixin{}.Fields()...)

	return fields
}

func (UniqueResourceMixin) Indexes() []ent.Index {
	indexes := ResourceMixin{}.Indexes()

	// Key mixin indexes are not used, as now that we know we have namespaces, we can use a better index

	// Soft deleted items should not create a conflict with a new item with the same key.
	// The proper index would be:
	//
	// 	CREATE UNIQE INDEX x ON y (namespace, key) WHERE deleted_at IS NULL
	//
	// ENT only supports WHERE clauses on indexes via manually creating migrations, so
	// we could approximate that behavior using this index.
	//
	// Caveats: If two resources with the same key are deleted in the same microsecond then the
	// deletion will fail. (e.g. by doing a create, delete, create, delete in the same microsecond)
	indexes = append(indexes, index.Fields("namespace", "key", "deleted_at").Unique())

	return indexes
}

// ResourceMixin adds common fields
type ResourceMixin struct {
	mixin.Schema
	MetadataDeprecatedReason string
}

func (m ResourceMixin) Fields() []ent.Field {
	var fields []ent.Field
	fields = append(fields, IDMixin{}.Fields()...)
	fields = append(fields, NamespaceMixin{}.Fields()...)
	fields = append(fields, MetadataMixin{DeprecatedReason: m.MetadataDeprecatedReason}.Fields()...)
	fields = append(fields, TimeMixin{}.Fields()...)
	fields = append(fields,
		field.String("name"),
		field.String("description").Optional().Nillable(),
	)

	return fields
}

func (ResourceMixin) Indexes() []ent.Index {
	var indexes []ent.Index
	indexes = append(indexes, IDMixin{}.Indexes()...)
	indexes = append(indexes, NamespaceMixin{}.Indexes()...)
	indexes = append(indexes, MetadataMixin{}.Indexes()...)
	indexes = append(indexes, TimeMixin{}.Indexes()...)
	indexes = append(indexes, index.Fields("namespace", "id").Unique())

	return indexes
}

// IDMixin adds the ID field to the schema
type IDMixin struct {
	mixin.Schema
}

// Fields of the IDMixin.
func (IDMixin) Fields() []ent.Field {
	return []ent.Field{
		field.String("id").
			DefaultFunc(func() string {
				return ulid.Make().String()
			}).
			Unique().
			Immutable().
			SchemaType(map[string]string{
				dialect.Postgres: "char(26)",
			}),
	}
}

func (IDMixin) Indexes() []ent.Index {
	return []ent.Index{
		index.Fields("id").Unique(),
	}
}

type IDMixinGetter interface {
	GetID() string
}

type IDMixinCreator[T any] interface {
	SetID(id string) T
}

// KeyMixin adds the key field to the schema
type KeyMixin struct {
	mixin.Schema
}

// Fields of the KeyMixin.
func (KeyMixin) Fields() []ent.Field {
	return []ent.Field{
		field.String("key").
			NotEmpty().
			Immutable(),
	}
}

// NamespaceMixin can be used for namespaced entities
type NamespaceMixin struct {
	mixin.Schema
}

// Fields of the IDMixin.
func (NamespaceMixin) Fields() []ent.Field {
	return []ent.Field{
		field.String("namespace").
			NotEmpty().
			Immutable(),
	}
}

func (NamespaceMixin) Indexes() []ent.Index {
	return []ent.Index{
		index.Fields("namespace"),
	}
}

type NamespaceMixinGetter interface {
	GetNamespace() string
}

type NamespaceMixinCreator[T any] interface {
	SetNamespace(namespace string) T
}

// MetadataMixin adds metadata to the schema
type MetadataMixin struct {
	mixin.Schema
	DeprecatedReason string
}

// Fields of the IDMixin.
func (m MetadataMixin) Fields() []ent.Field {
	metadata := field.JSON("metadata", map[string]string{}).
		Optional().
		SchemaType(map[string]string{
			dialect.Postgres: "jsonb",
		})
	if m.DeprecatedReason != "" {
		metadata.Deprecated(m.DeprecatedReason)
	}

	return []ent.Field{metadata}
}

// AnnotationsMixin adds annotations to the schema
type AnnotationsMixin struct {
	mixin.Schema
	DeprecatedReason string
}

// Fields of the IDMixin.
func (m AnnotationsMixin) Fields() []ent.Field {
	annotations := field.JSON("annotations", models.Annotations{}).
		Optional().
		SchemaType(map[string]string{
			dialect.Postgres: "jsonb",
		})
	if m.DeprecatedReason != "" {
		annotations.Deprecated(m.DeprecatedReason)
	}

	return []ent.Field{annotations}
}

func (AnnotationsMixin) Indexes() []ent.Index {
	return []ent.Index{
		index.Fields("annotations").
			Annotations(
				entsql.IndexTypes(map[string]string{
					dialect.Postgres: "GIN",
				}),
			),
	}
}

type AnnotationsMixinGetter interface {
	GetAnnotations() models.Annotations
}

type AnnotationsMixinSetter[T any] interface {
	SetAnnotations(annotations models.Annotations) T
}

// TimeMixin adds the created_at and updated_at fields to the schema
type TimeMixin struct {
	mixin.Schema
}

// Fields of the TimeMixin.
func (TimeMixin) Fields() []ent.Field {
	return []ent.Field{
		field.Time("created_at").
			Default(truncatedNow).
			Immutable(),
		field.Time("updated_at").
			Default(truncatedNow).
			UpdateDefault(truncatedNow),
		field.Time("deleted_at").
			Optional().
			Nillable(),
	}
}

type TimeMixinGetter interface {
	GetCreatedAt() time.Time
	GetUpdatedAt() time.Time
	GetDeletedAt() *time.Time
}

type TimeMixinCreator[T any] interface {
	SetCreatedAt(createdAt time.Time) T
	SetUpdatedAt(updatedAt time.Time) T
	SetNillableDeletedAt(deletedAt *time.Time) T
}

type TimeMixinUpdater[T any] interface {
	SetUpdatedAt(updatedAt time.Time) T
	SetNillableDeletedAt(deletedAt *time.Time) T
}

func MapTimeMixinFromDB[T TimeMixinGetter](dbEntity T) models.ManagedModel {
	return models.ManagedModel{
		CreatedAt: dbEntity.GetCreatedAt().In(time.UTC),
		UpdatedAt: dbEntity.GetUpdatedAt().In(time.UTC),
		DeletedAt: convert.TimePtrIn(dbEntity.GetDeletedAt(), time.UTC),
	}
}

// truncatedNow returns the current time truncated to microsecond precision. This is useful, as:
// - ent when creating resources will return the in memory calculated data including the timestamp in host precision
// - PostgreSQL has microsecond precision
// - Linux has nanosecond precision
// - MacOS seem to have at most microsecond precision in go
//
// This means that any test that relies on CreatedAt or UpdatedAt comparisons will pass on macos, but will fail on CI.
func truncatedNow() time.Time {
	// PostgreSQL has microsecond precision, so let's truncate to that which makes
	// it easier to test and compare times.
	return clock.Now().Truncate(time.Microsecond)
}

type CadencedMixin struct {
	mixin.Schema
}

func (CadencedMixin) Fields() []ent.Field {
	return []ent.Field{
		field.Time("active_from").Immutable(),
		field.Time("active_to").Optional().Nillable(),
	}
}

// CustomerAddressMixin adds address fields to a customer, used by billing to snapshot addresses for invoices
type CustomerAddressMixin struct {
	ent.Schema
	FieldPrefix string
}

func (c CustomerAddressMixin) Fields() []ent.Field {
	return []ent.Field{
		// PII fields
		field.String(fmt.Sprintf("%s_address_country", c.FieldPrefix)).GoType(models.CountryCode("")).MinLen(2).MaxLen(2).Optional().Nillable(),
		field.String(fmt.Sprintf("%s_address_postal_code", c.FieldPrefix)).Optional().Nillable(),
		field.String(fmt.Sprintf("%s_address_state", c.FieldPrefix)).Optional().Nillable(),
		field.String(fmt.Sprintf("%s_address_city", c.FieldPrefix)).Optional().Nillable(),
		field.String(fmt.Sprintf("%s_address_line1", c.FieldPrefix)).Optional().Nillable(),
		field.String(fmt.Sprintf("%s_address_line2", c.FieldPrefix)).Optional().Nillable(),
		field.String(fmt.Sprintf("%s_address_phone_number", c.FieldPrefix)).Optional().Nillable(),
	}
}