executor.go

package fscli

import (
	"context"
	"errors"

	"cloud.google.com/go/firestore"
	"cloud.google.com/go/firestore/apiv1/firestorepb"
	"google.golang.org/api/iterator"
)

type Executor struct {
	fs *firestore.Client
}

var ErrInvalidCollection = errors.New("invalid collection")

func NewExecutor(ctx context.Context, fs *firestore.Client) *Executor {
	return &Executor{fs}
}

func (exe *Executor) ExecuteQuery(ctx context.Context, op *QueryOperation) ([]*firestore.DocumentSnapshot, error) {
	var q firestore.Query
	if op.IsCollectionGroup() {
		q = exe.fs.CollectionGroup(op.Collection()).Query
	} else {
		collection := exe.fs.Collection(op.Collection())
		if collection == nil {
			return nil, ErrInvalidCollection
		}
		q = collection.Query
	}

	for _, filter := range op.filters {
		fieldName := filter.FieldName()
		value := filter.Value()
		if fieldName == FieldDocumentID {
			fieldName = firestore.DocumentID
			if !op.IsCollectionGroup() {
				value = toDocRefValue(exe.fs.Collection(op.Collection()), value)
			}
		}
		q = q.Where(fieldName, string(filter.Operator()), value)
	}

	if len(op.selects) > 0 {
		q = q.Select(op.selects...)
	}

	if len(op.orderBys) > 0 {
		for _, orderBy := range op.orderBys {
			q = q.OrderBy(orderBy.field, firestore.Direction(orderBy.direction))
		}
	}

	if op.limit > 0 {
		q = q.Limit(op.limit)
	}

	itr := q.Documents(ctx)
	defer itr.Stop()

	docs := make([]*firestore.DocumentSnapshot, 0)
	for {
		doc, err := itr.Next()
		if err == iterator.Done {
			break
		}
		if err != nil {
			return nil, err
		}

		docs = append(docs, doc)
	}
	return docs, nil
}

func (exe *Executor) ExecuteGet(ctx context.Context, op *GetOperation) (*firestore.DocumentSnapshot, error) {
	doc, err := exe.fs.Collection(op.Collection()).Doc(op.DocId()).Get(ctx)
	if err != nil {
		return nil, err
	}
	return doc, nil
}

func (exe *Executor) ExecuteCount(ctx context.Context, op *CountOperation) (int64, error) {
	var q firestore.Query
	if op.IsCollectionGroup() {
		q = exe.fs.CollectionGroup(op.Collection()).Query
	} else {
		collection := exe.fs.Collection(op.Collection())
		if collection == nil {
			return 0, ErrInvalidCollection
		}
		q = collection.Query
	}

	for _, filter := range op.filters {
		fieldName := filter.FieldName()
		value := filter.Value()
		if fieldName == FieldDocumentID {
			fieldName = firestore.DocumentID
			if !op.IsCollectionGroup() {
				value = toDocRefValue(exe.fs.Collection(op.Collection()), value)
			}
		}
		q = q.Where(fieldName, string(filter.Operator()), value)
	}

	aggrQ := q.NewAggregationQuery().WithCount("all")
	results, err := aggrQ.Get(ctx)
	if err != nil {
		return 0, err
	}
	count, ok := results["all"]
	if !ok {
		return 0, errors.New("invalid aggregation result")
	}

	v := count.(*firestorepb.Value)
	return v.GetIntegerValue(), nil
}

func (exe *Executor) ExecuteListCollections(ctx context.Context, cmd *MetacommandListCollections) ([]string, error) {
	return findAllCollections(ctx, exe.fs, cmd.baseDoc)
}

func toDocRefValue(collection *firestore.CollectionRef, value any) any {
	switch v := value.(type) {
	case string:
		return collection.Doc(v)
	case []any:
		refs := make([]*firestore.DocumentRef, len(v))
		for i, item := range v {
			if s, ok := item.(string); ok {
				refs[i] = collection.Doc(s)
			}
		}
		return refs
	default:
		return value
	}
}