signoz/pkg/sqlschema/index.go

package sqlschema

import (
	"fmt"
	"hash/fnv"
	"slices"
	"strings"

	"github.com/SigNoz/signoz/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
}

// UniqueIndexWithExpressions is a functional unique index: each key is a SQL
// expression (e.g. LOWER(col)) emitted verbatim, so the caller owns its
// well-formedness; plain columns go in as bare identifiers. The auto-generated
// name uses a hash suffix (uq_t_<hash>) since expressions aren't valid
// identifier fragments.
//
// Use this only when at least one key is a real expression: an all-identifier
// key list reconstructs as a plain UniqueIndex on read-back and won't compare
// equal — use UniqueIndex instead.
type UniqueIndexWithExpressions struct {
	TableName   TableName
	Expressions []string
	name        string
}

func (index *UniqueIndexWithExpressions) 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("_")
	hasher := fnv.New32a()
	_, _ = hasher.Write([]byte(strings.Join(normalizeExpressions(index.Expressions), "\x00")))
	fmt.Fprintf(&b, "%08x", hasher.Sum32())
	return b.String()
}

func (index *UniqueIndexWithExpressions) Named(name string) Index {
	copyOfExpressions := make([]string, len(index.Expressions))
	copy(copyOfExpressions, index.Expressions)

	return &UniqueIndexWithExpressions{
		TableName:   index.TableName,
		Expressions: copyOfExpressions,
		name:        name,
	}
}

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

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

// Columns is nil: a functional index exposes no plain key columns.
func (*UniqueIndexWithExpressions) Columns() []ColumnName {
	return nil
}

func (index *UniqueIndexWithExpressions) Equals(other Index) bool {
	otherIndex, ok := other.(*UniqueIndexWithExpressions)
	if !ok {
		return false
	}

	// Expressions are compared normalized so a declared `LOWER(x)` matches the
	// `lower(x)` a backend renders on read-back.
	return index.Name() == other.Name() && slices.Equal(normalizeExpressions(index.Expressions), normalizeExpressions(otherIndex.Expressions))
}

// normalizeExpressions canonicalizes each Expressions entry (case, whitespace,
// redundant parens, identifier quoting) so a declared LOWER(x) matches the
// lower(x) a backend renders on read-back. Reuses expressionNormalizer: an index key
// is the same SQL grammar as a WHERE predicate. It does not handle opclass,
// COLLATE, or ASC/DESC/NULLS suffixes: postgres strips them on read-back
// (pg_get_indexdef), while sqlite keeps them verbatim, so an index declared with
// one round-trips on sqlite but not on postgres.
func normalizeExpressions(expressions []string) []string {
	if len(expressions) == 0 {
		return nil
	}

	normalized := make([]string, len(expressions))
	for i, expression := range expressions {
		normalized[i] = (&expressionNormalizer{input: expression, foldCase: true}).normalize()
	}
	return normalized
}

func (index *UniqueIndexWithExpressions) 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, expr := range index.Expressions {
		if i > 0 {
			sql = append(sql, ", "...)
		}
		sql = append(sql, expr...)
	}

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

	return sql
}

func (index *UniqueIndexWithExpressions) 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((&expressionNormalizer{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()) && (&expressionNormalizer{input: index.Where}).normalize() == (&expressionNormalizer{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
}