dynamo-browse/internal/common/ui/commandctrl/cmdpacks/modrs.go

package cmdpacks

import (
	"context"
	"github.com/aws/aws-sdk-go-v2/service/dynamodb/types"
	"lmika.dev/cmd/dynamo-browse/internal/common/ui/commandctrl"
	"lmika.dev/cmd/dynamo-browse/internal/dynamo-browse/controllers"
	"lmika.dev/cmd/dynamo-browse/internal/dynamo-browse/models"
	"lmika.dev/cmd/dynamo-browse/internal/dynamo-browse/models/queryexpr"
	"lmika.dev/cmd/dynamo-browse/internal/dynamo-browse/services/tables"
	"github.com/pkg/errors"
	"time"
	"ucl.lmika.dev/repl"
	"ucl.lmika.dev/ucl"
)

type rsModule struct {
	tableService *tables.Service
	state        *controllers.State
}

var rsNewDoc = repl.Doc{
	Brief: "Creates a new, empty result set",
	Usage: "[-table NAME]",
	Detailed: `
		The result set assumes the details of the current table. If no table is specified,
		the command will return an error.
	`,
}

func (rs *rsModule) rsNew(ctx context.Context, args ucl.CallArgs) (_ any, err error) {
	var tableInfo *models.TableInfo
	if args.HasSwitch("table") {
		var tblName string
		if err := args.BindSwitch("table", &tblName); err != nil {
			return nil, err
		}

		tableInfo, err = rs.tableService.Describe(ctx, tblName)
		if err != nil {
			return nil, err
		}
	} else if currRs := rs.state.ResultSet(); currRs != nil && currRs.TableInfo != nil {
		tableInfo = currRs.TableInfo
	} else {
		return nil, errors.New("no table specified")
	}

	return newResultSetProxy(&models.ResultSet{
		TableInfo: tableInfo,
		Created:   time.Now(),
	}), nil
}

var rsQueryDoc = repl.Doc{
	Brief: "Runs a query and returns the results as a result-set",
	Usage: "QUERY [ARGS] [-table NAME]",
	Args: []repl.ArgDoc{
		{Name: "query", Brief: "Query expression to run"},
		{Name: "args", Brief: "Hash of argument values to substitute into the query"},
		{Name: "-table", Brief: "Optional table name to use for the query"},
	},
	Detailed: `
		If no table is specified, then the value of @table will be used. If this is unavailable,
		the command will return an error.
    `,
}

func parseQuery(
	ctx context.Context,
	args ucl.CallArgs,
	currentRS *models.ResultSet,
	tablesService *tables.Service,
) (*queryexpr.QueryExpr, *models.TableInfo, error) {
	var expr string
	if err := args.Bind(&expr); err != nil {
		return nil, nil, err
	}

	q, err := queryexpr.Parse(expr)
	if err != nil {
		return nil, nil, err
	}

	if args.NArgs() > 0 {
		var queryArgs ucl.Hashable
		if err := args.Bind(&queryArgs); err != nil {
			return nil, nil, err
		}

		queryNames := map[string]string{}
		queryValues := map[string]types.AttributeValue{}
		queryArgs.Each(func(k string, v ucl.Object) error {
			if v == nil {
				return nil
			}

			queryNames[k] = v.String()

			switch v.(type) {
			case ucl.StringObject:
				queryValues[k] = &types.AttributeValueMemberS{Value: v.String()}
			case ucl.IntObject:
				queryValues[k] = &types.AttributeValueMemberN{Value: v.String()}
				// TODO: other types
			}
			return nil
		})

		q = q.WithNameParams(queryNames).WithValueParams(queryValues)
	}

	var tableInfo *models.TableInfo
	if args.HasSwitch("table") {
		var tblName string
		if err := args.BindSwitch("table", &tblName); err != nil {
			return nil, nil, err
		}

		tableInfo, err = tablesService.Describe(ctx, tblName)
		if err != nil {
			return nil, nil, err
		}
	} else if currentRS != nil && currentRS.TableInfo != nil {
		tableInfo = currentRS.TableInfo
	} else {
		return nil, nil, errors.New("no table specified")
	}

	return q, tableInfo, nil
}

func (rs *rsModule) rsQuery(ctx context.Context, args ucl.CallArgs) (any, error) {
	q, tableInfo, err := parseQuery(ctx, args, rs.state.ResultSet(), rs.tableService)
	if err != nil {
		return nil, err
	}

	newResultSet, err := rs.tableService.ScanOrQuery(context.Background(), tableInfo, q, nil)
	if err != nil {
		return nil, err
	}

	return newResultSetProxy(newResultSet), nil
}

var rsScanDoc = repl.Doc{
	Brief: "Performs a scan of the table and returns the results as a result-set",
	Usage: "[-table NAME]",
	Args: []repl.ArgDoc{
		{Name: "-table", Brief: "Optional table name to use for the query"},
	},
	Detailed: `
		If no table is specified, then the value of @table will be used. If this is unavailable,
		the command will return an error.
    `,
}

func (rs *rsModule) rsScan(ctx context.Context, args ucl.CallArgs) (_ any, err error) {
	var tableInfo *models.TableInfo
	if args.HasSwitch("table") {
		var tblName string
		if err := args.BindSwitch("table", &tblName); err != nil {
			return nil, err
		}

		tableInfo, err = rs.tableService.Describe(ctx, tblName)
		if err != nil {
			return nil, err
		}
	} else if currRs := rs.state.ResultSet(); currRs != nil && currRs.TableInfo != nil {
		tableInfo = currRs.TableInfo
	} else {
		return nil, errors.New("no table specified")
	}

	newResultSet, err := rs.tableService.Scan(context.Background(), tableInfo)
	if err != nil {
		return nil, err
	}

	return newResultSetProxy(newResultSet), nil
}

func (rs *rsModule) rsFilter(ctx context.Context, args ucl.CallArgs) (any, error) {
	var (
		rsProxy SimpleProxy[*models.ResultSet]
		filter  string
	)

	if err := args.Bind(&rsProxy, &filter); err != nil {
		return nil, err
	}

	newResultSet := rs.tableService.Filter(rsProxy.ProxyValue(), filter)
	return newResultSetProxy(newResultSet), nil
}

var rsNextPageDoc = repl.Doc{
	Brief: "Returns the next page of the passed in result-set",
	Usage: "RESULT_SET",
	Args: []repl.ArgDoc{
		{Name: "result-set", Brief: "Result set to fetch the next page of"},
	},
	Detailed: `
		If no next page exists, the command will return nil.
	`,
}

func (rs *rsModule) rsNextPage(ctx context.Context, args ucl.CallArgs) (_ any, err error) {
	var rsProxy SimpleProxy[*models.ResultSet]

	if err := args.Bind(&rsProxy); err != nil {
		return nil, err
	}

	if !rsProxy.value.HasNextPage() {
		return nil, nil
	}

	nextPage, err := rs.tableService.NextPage(ctx, rsProxy.value)
	if err != nil {
		return nil, err
	}

	return newResultSetProxy(nextPage), nil
}

func (rs *rsModule) rsUnion(ctx context.Context, args ucl.CallArgs) (_ any, err error) {
	var rsProxy1, rsProxy2 SimpleProxy[*models.ResultSet]

	if err := args.Bind(&rsProxy1, &rsProxy2); err != nil {
		return nil, err
	}

	return newResultSetProxy(rsProxy1.ProxyValue().MergeWith(rsProxy2.ProxyValue())), nil
}

func (rs *rsModule) rsSet(ctx context.Context, args ucl.CallArgs) (_ any, err error) {
	var (
		item itemProxy
		expr string
		val  ucl.Object
	)

	if err := args.Bind(&item, &expr, &val); err != nil {
		return nil, err
	}

	q, err := queryexpr.Parse(expr)
	if err != nil {
		return nil, err
	}

	// TEMP: attribute is always S
	if err := q.SetEvalItem(item.item, &types.AttributeValueMemberS{Value: val.String()}); err != nil {
		return nil, err
	}
	item.resultSet.SetDirty(item.idx, true)
	commandctrl.QueueRefresh(ctx)

	return item, nil
}

func (rs *rsModule) rsDel(ctx context.Context, args ucl.CallArgs) (_ any, err error) {
	var (
		item itemProxy
		expr string
	)

	if err := args.Bind(&item, &expr); err != nil {
		return nil, err
	}

	q, err := queryexpr.Parse(expr)
	if err != nil {
		return nil, err
	}

	if err := q.DeleteAttribute(item.item); err != nil {
		return nil, err
	}
	item.resultSet.SetDirty(item.idx, true)
	commandctrl.QueueRefresh(ctx)

	return item, nil
}

func moduleRS(tableService *tables.Service, state *controllers.State) ucl.Module {
	m := &rsModule{
		tableService: tableService,
		state:        state,
	}

	return ucl.Module{
		Name: "rs",
		Builtins: map[string]ucl.BuiltinHandler{
			"new":       m.rsNew,
			"query":     m.rsQuery,
			"scan":      m.rsScan,
			"filter":    m.rsFilter,
			"next-page": m.rsNextPage,
			"union":     m.rsUnion,
			"set":       m.rsSet,
			"del":       m.rsDel,
		},
	}
}