chaisql/internal/stream/index/insert.go

package index

import (
	"fmt"

	"github.com/chaisql/chai/internal/environment"
	"github.com/chaisql/chai/internal/stream"
	"github.com/chaisql/chai/internal/types"
	"github.com/cockroachdb/errors"
)

// InsertOperator reads the input stream and indexes each row.
type InsertOperator struct {
	stream.BaseOperator

	indexName string
}

func Insert(indexName string) *InsertOperator {
	return &InsertOperator{
		indexName: indexName,
	}
}

func (op *InsertOperator) Clone() stream.Operator {
	return &InsertOperator{
		BaseOperator: op.BaseOperator.Clone(),
		indexName:    op.indexName,
	}
}

func (op *InsertOperator) Iterate(in *environment.Environment, fn func(out *environment.Environment) error) error {
	tx := in.GetTx()

	idx, err := tx.Catalog.GetIndex(tx, op.indexName)
	if err != nil {
		return err
	}

	info, err := tx.Catalog.GetIndexInfo(op.indexName)
	if err != nil {
		return err
	}

	tinfo, err := tx.Catalog.GetTableInfo(info.Owner.TableName)
	if err != nil {
		return err
	}

	return op.Prev.Iterate(in, func(out *environment.Environment) error {
		r, ok := out.GetDatabaseRow()
		if !ok {
			return errors.New("missing row")
		}

		vs := make([]types.Value, 0, len(info.Columns))
		for _, column := range info.Columns {
			v, err := r.Get(column)
			if err != nil {
				v = types.NewNullValue()
			}
			vs = append(vs, v)
		}

		encKey, err := tinfo.EncodeKey(r.Key())
		if err != nil {
			return err
		}

		err = idx.Set(vs, encKey)
		if err != nil {
			return fmt.Errorf("error while inserting index value: %w", err)
		}

		return fn(out)
	})
}

func (op *InsertOperator) String() string {
	return fmt.Sprintf("index.Insert(%q)", op.indexName)
}