upstream-ory-keto/internal/persistence/sql/uuid_mapping.go

// Copyright © 2023 Ory Corp
// SPDX-License-Identifier: Apache-2.0

package sql

import (
	"bytes"
	"context"
	"iter"
	"maps"
	"slices"
	"strings"

	"github.com/gofrs/uuid"
	"github.com/ory/x/otelx"
	"github.com/ory/x/pagination/keysetpagination"
	"github.com/ory/x/sqlcon"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/trace"
)

type (
	UUIDMapping struct {
		ID                   uuid.UUID `db:"id"`
		StringRepresentation string    `db:"string_representation"`
	}
	UUIDMappings []*UUIDMapping
)

func (UUIDMappings) TableName() string {
	return "keto_uuid_mappings"
}

func (UUIDMapping) TableName() string {
	return "keto_uuid_mappings"
}

func (p *Persister) batchFromUUIDs(ctx context.Context, ids []uuid.UUID) (res []string, err error) {
	if len(ids) == 0 {
		return
	}

	p.d.Logger().Trace("looking up UUIDs")

	// Build a map from UUID -> indices in the result.
	idIdx := make(map[uuid.UUID][]int, len(ids))
	for i, id := range ids {
		if ids, ok := idIdx[id]; ok {
			idIdx[id] = append(ids, i)
		} else {
			idIdx[id] = []int{i}
		}
	}
	nextID, stop := iter.Pull(maps.Keys(idIdx))
	defer stop()

	res = make([]string, len(ids))

	idsToLookup := make([]uuid.UUID, 0, keysetpagination.DefaultMaxSize)
	mappings := make([]UUIDMapping, 0, keysetpagination.DefaultMaxSize)
	for i := 0; i < len(idIdx); i += keysetpagination.DefaultMaxSize {
		idsToLookup = idsToLookup[:0]
		mappings = mappings[:0]

		for range keysetpagination.DefaultMaxSize {
			id, ok := nextID()
			if !ok {
				break
			}
			idsToLookup = append(idsToLookup, id)
		}
		query := p.Connection(ctx).Where("id in (?)", idsToLookup)
		if err := sqlcon.HandleError(query.All(&mappings)); err != nil {
			return []string{}, err
		}

		// Write the representation to the correct index.
		for _, m := range mappings {
			for _, idx := range idIdx[m.ID] {
				res[idx] = m.StringRepresentation
			}
		}
	}

	return
}

func (p *Persister) MapStringsToUUIDs(ctx context.Context, values ...string) (uuids []uuid.UUID, err error) {
	if len(values) == 0 {
		return
	}

	ctx, span := p.d.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.MapStringsToUUIDs",
		trace.WithAttributes(attribute.Int("num_values", len(values))))
	defer otelx.End(span, &err)

	uuids, err = p.MapStringsToUUIDsReadOnly(ctx, values...)
	if err != nil {
		return nil, err
	}

	p.d.Logger().WithField("values", values).WithField("UUIDs", uuids).Trace("adding UUID mappings")

	mappings := make([]UUIDMapping, len(values))
	for i := range values {
		mappings[i] = UUIDMapping{
			ID:                   uuids[i],
			StringRepresentation: values[i],
		}
	}
	slices.SortFunc(mappings, func(a, b UUIDMapping) int {
		return bytes.Compare(a.ID[:], b.ID[:])
	})
	mappings = slices.CompactFunc(mappings, func(a, b UUIDMapping) bool {
		return a.ID == b.ID
	})

	span.SetAttributes(attribute.Int("num_mappings", len(mappings)))

	err = p.Transaction(ctx, func(ctx context.Context) error {
		for chunk := range slices.Chunk(mappings, chunkSizeInsertUUIDMappings) {
			query, args := buildInsertUUIDs(chunk, p.conn.Dialect.Name())
			if err := p.Connection(ctx).RawQuery(query, args...).Exec(); err != nil {
				return sqlcon.HandleError(err)
			}
		}
		return nil
	})

	return uuids, err
}

func (p *Persister) MapStringsToUUIDsReadOnly(ctx context.Context, ss ...string) (uuids []uuid.UUID, err error) {
	// This function doesn't talk to the database or do anything interesting, so we don't need to trace it.
	// ctx, span := p.d.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.MapStringsToUUIDsReadOnly")
	// defer otelx.End(span, &err)

	uuids = make([]uuid.UUID, len(ss))
	for i := range ss {
		uuids[i] = uuid.NewV5(p.d.MapperNamespace(ctx), ss[i])
	}
	return uuids, nil
}

func (p *Persister) MapUUIDsToStrings(ctx context.Context, u ...uuid.UUID) (_ []string, err error) {
	ctx, span := p.d.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.MapUUIDsToStrings")
	defer otelx.End(span, &err)

	return p.batchFromUUIDs(ctx, u)
}

func buildInsertUUIDs(values []UUIDMapping, dialect string) (query string, args []any) {
	if len(values) == 0 {
		return "", nil
	}

	const placeholder = "(?,?)"
	const separator = ","

	var q strings.Builder
	args = make([]any, 0, len(values)*2)

	if dialect == "mysql" {
		q.WriteString("INSERT IGNORE INTO keto_uuid_mappings (id, string_representation) VALUES ")
	} else {
		q.WriteString("INSERT INTO keto_uuid_mappings (id, string_representation) VALUES ")
	}

	q.Grow(len(values)*(len(placeholder)+len(separator)) + 100)

	for i, val := range values {
		if i > 0 {
			q.WriteString(separator)
		}
		q.WriteString(placeholder)
		args = append(args, val.ID, val.StringRepresentation)
	}

	if dialect == "mysql" {
		// nothing
	} else {
		q.WriteString(" ON CONFLICT (id) DO NOTHING")
	}

	return q.String(), args
}