o11y/pkg/sqlschema/index.go

package sqlschema

import (
	"slices"
	"strings"

	"github.com/hanzoai/o11y/pkg/valuer"
)

var (
	IndexTypeUnique        = IndexType{s: valuer.NewString("uq")}
	IndexTypeIndex         = IndexType{s: valuer.NewString("ix")}
	IndexTypePartialUnique = IndexType{s: valuer.NewString("puq")}
)

type IndexType struct{ s valuer.String }

func (i IndexType) String() string {
	return i.s.String()
}

type Index interface {
	// The name of the index.
	//   - Indexes are named as `ix_<table_name>_<column_names>`. The column names are separated by underscores.
	//   - Unique constraints are named as `uq_<table_name>_<column_names>`. The column names are separated by underscores.
	//   - Partial unique indexes are named as `puq_<table_name>_<column_names>_<predicate_hash>`.
	// The name is autogenerated and should not be set by the user.
	Name() string

	// Add name to the index. This is typically used to override the autogenerated name because the database might have a different name.
	Named(name string) Index

	// Returns true if the index is named. A named index is not autogenerated
	IsNamed() bool

	// The type of the index.
	Type() IndexType

	// The columns that the index is applied to.
	Columns() []ColumnName

	// Equals returns true if the index is equal to the other index.
	Equals(other Index) bool

	// The SQL representation of the index.
	ToCreateSQL(fmter SQLFormatter) []byte

	// Drop the index.
	ToDropSQL(fmter SQLFormatter) []byte
}

type UniqueIndex struct {
	TableName   TableName
	ColumnNames []ColumnName
	name        string
}

func (index *UniqueIndex) Name() string {
	if index.name != "" {
		return index.name
	}

	var b strings.Builder
	b.WriteString(IndexTypeUnique.String())
	b.WriteString("_")
	b.WriteString(string(index.TableName))
	b.WriteString("_")
	for i, column := range index.ColumnNames {
		if i > 0 {
			b.WriteString("_")
		}
		b.WriteString(string(column))
	}
	return b.String()
}

func (index *UniqueIndex) Named(name string) Index {
	copyOfColumnNames := make([]ColumnName, len(index.ColumnNames))
	copy(copyOfColumnNames, index.ColumnNames)

	return &UniqueIndex{
		TableName:   index.TableName,
		ColumnNames: copyOfColumnNames,
		name:        name,
	}
}

func (index *UniqueIndex) IsNamed() bool {
	return index.name != ""
}

func (*UniqueIndex) Type() IndexType {
	return IndexTypeUnique
}

func (index *UniqueIndex) Columns() []ColumnName {
	return index.ColumnNames
}

func (index *UniqueIndex) Equals(other Index) bool {
	if other.Type() != IndexTypeUnique {
		return false
	}

	return index.Name() == other.Name() && slices.Equal(index.Columns(), other.Columns())
}

func (index *UniqueIndex) ToCreateSQL(fmter SQLFormatter) []byte {
	sql := []byte{}

	sql = append(sql, "CREATE UNIQUE INDEX IF NOT EXISTS "...)
	sql = fmter.AppendIdent(sql, index.Name())
	sql = append(sql, " ON "...)
	sql = fmter.AppendIdent(sql, string(index.TableName))
	sql = append(sql, " ("...)

	for i, column := range index.ColumnNames {
		if i > 0 {
			sql = append(sql, ", "...)
		}

		sql = fmter.AppendIdent(sql, string(column))
	}

	sql = append(sql, ")"...)

	return sql
}

func (index *UniqueIndex) ToDropSQL(fmter SQLFormatter) []byte {
	sql := []byte{}

	sql = append(sql, "DROP INDEX IF EXISTS "...)
	sql = fmter.AppendIdent(sql, index.Name())

	return sql
}

type PartialUniqueIndex struct {
	TableName   TableName
	ColumnNames []ColumnName
	Where       string
	name        string
}

func (index *PartialUniqueIndex) Name() string {
	if index.name != "" {
		return index.name
	}

	var b strings.Builder
	b.WriteString(IndexTypePartialUnique.String())
	b.WriteString("_")
	b.WriteString(string(index.TableName))
	b.WriteString("_")
	for i, column := range index.ColumnNames {
		if i > 0 {
			b.WriteString("_")
		}
		b.WriteString(string(column))
	}
	b.WriteString("_")
	b.WriteString((&whereNormalizer{input: index.Where}).hash())
	return b.String()
}

func (index *PartialUniqueIndex) Named(name string) Index {
	copyOfColumnNames := make([]ColumnName, len(index.ColumnNames))
	copy(copyOfColumnNames, index.ColumnNames)

	return &PartialUniqueIndex{
		TableName:   index.TableName,
		ColumnNames: copyOfColumnNames,
		Where:       index.Where,
		name:        name,
	}
}

func (index *PartialUniqueIndex) IsNamed() bool {
	return index.name != ""
}

func (*PartialUniqueIndex) Type() IndexType {
	return IndexTypePartialUnique
}

func (index *PartialUniqueIndex) Columns() []ColumnName {
	return index.ColumnNames
}

func (index *PartialUniqueIndex) Equals(other Index) bool {
	if other.Type() != IndexTypePartialUnique {
		return false
	}

	otherPartial, ok := other.(*PartialUniqueIndex)
	if !ok {
		return false
	}

	return index.Name() == other.Name() && slices.Equal(index.Columns(), other.Columns()) && (&whereNormalizer{input: index.Where}).normalize() == (&whereNormalizer{input: otherPartial.Where}).normalize()
}

func (index *PartialUniqueIndex) ToCreateSQL(fmter SQLFormatter) []byte {
	sql := []byte{}

	sql = append(sql, "CREATE UNIQUE INDEX IF NOT EXISTS "...)
	sql = fmter.AppendIdent(sql, index.Name())
	sql = append(sql, " ON "...)
	sql = fmter.AppendIdent(sql, string(index.TableName))
	sql = append(sql, " ("...)

	for i, column := range index.ColumnNames {
		if i > 0 {
			sql = append(sql, ", "...)
		}

		sql = fmter.AppendIdent(sql, string(column))
	}

	sql = append(sql, ") WHERE "...)
	sql = append(sql, index.Where...)

	return sql
}

func (index *PartialUniqueIndex) ToDropSQL(fmter SQLFormatter) []byte {
	sql := []byte{}

	sql = append(sql, "DROP INDEX IF EXISTS "...)
	sql = fmter.AppendIdent(sql, index.Name())

	return sql
}