[raft/uss] Extract uss availability

pkg/scd/actions/availability.go

@@ -0,0 +1,104 @@
+package actions
+
+import (
+	"context"
+
+	restapi "github.com/interuss/dss/pkg/api/scdv1"
+	dsserr "github.com/interuss/dss/pkg/errors"
+	dssmodels "github.com/interuss/dss/pkg/models"
+	scdmodels "github.com/interuss/dss/pkg/scd/models"
+	"github.com/interuss/dss/pkg/scd/repos"
+	dssstore "github.com/interuss/dss/pkg/store"
+	"github.com/interuss/stacktrace"
+	"github.com/jackc/pgx/v5"
+)
+
+func init() {
+	Registry[restapi.GetUssAvailabilityOperationID] = dssstore.OperationHandler[repos.Repository]{
+		Encode:     dssstore.EncodeJSON,
+		Decode:     dssstore.DecodeJSON[*restapi.GetUssAvailabilityRequest],
+		Execute:    ExecuteGetUssAvailability,
+		IsReadOnly: true,
+	}
+	Registry[restapi.SetUssAvailabilityOperationID] = dssstore.OperationHandler[repos.Repository]{
+		Encode:  dssstore.EncodeJSON,
+		Decode:  dssstore.DecodeJSON[*restapi.SetUssAvailabilityRequest],
+		Execute: ExecuteSetUssAvailability,
+	}
+}
+
+func ExecuteGetUssAvailability(ctx context.Context, repo repos.Repository, request dssstore.OperationRequest) (any, error) {
+	req, ok := request.(*restapi.GetUssAvailabilityRequest)
+	if !ok {
+		return nil, stacktrace.NewError("unexpected request type %T for operation %q", request, restapi.GetUssAvailabilityOperationID)
+	}
+
+	id := dssmodels.ManagerFromString(req.UssId)
+
+	// Get USS availability from Store
+	ussa, err := repo.GetUssAvailability(ctx, id)
+	if err != nil && err != pgx.ErrNoRows {
+		return nil, stacktrace.Propagate(err, "Could not get USS availability from repo")
+	}
+	if ussa == nil {
+		// Return default availability status "Unknown"
+		return &restapi.UssAvailabilityStatusResponse{
+			Status: restapi.UssAvailabilityStatus{
+				Availability: restapi.UssAvailabilityState_Unknown,
+				Uss:          id.String()},
+		}, nil
+	}
+
+	return &restapi.UssAvailabilityStatusResponse{
+		Status:  *ussa.ToRest(),
+		Version: ussa.Version.String(),
+	}, nil
+}
+
+func ExecuteSetUssAvailability(ctx context.Context, repo repos.Repository, request dssstore.OperationRequest) (any, error) {
+	req, ok := request.(*restapi.SetUssAvailabilityRequest)
+	if !ok {
+		return nil, stacktrace.NewError("unexpected request type %T for operation %q", request, restapi.SetUssAvailabilityOperationID)
+	}
+
+	// Retrieve USS availability status from request params
+	availability, err := scdmodels.UssAvailabilityStateFromRest(req.Body.Availability)
+	if err != nil {
+		return nil, stacktrace.PropagateWithCode(err, dsserr.BadRequest, "Invalid availability state")
+	}
+	id := dssmodels.ManagerFromString(req.UssId)
+	version := scdmodels.OVN(req.Body.OldVersion)
+	ussareq := &scdmodels.UssAvailabilityStatus{
+		Uss:          id,
+		Availability: availability,
+	}
+
+	old, err := repo.GetUssAvailability(ctx, id)
+	if err != nil && err != pgx.ErrNoRows {
+		return nil, stacktrace.Propagate(err, "Could not get USS availability from repo")
+	}
+	switch {
+	case old == nil && !version.Empty():
+		// The user wants set a new availability status but it already exists.
+		return nil, stacktrace.NewErrorWithCode(dsserr.AlreadyExists, "availability for USS %s already exists", id.String())
+	case old != nil && old.Version != version:
+		// The user wants to update an availability status but the version doesn't match.
+		return nil, stacktrace.Propagate(
+			stacktrace.NewErrorWithCode(dsserr.VersionMismatch, "USS availability version %s is not current", version),
+			"Current version is %s but client specified version %s", old.Version, version)
+	}
+
+	// Upsert the USS availability
+	ussa, err := repo.UpsertUssAvailability(ctx, ussareq)
+	if err != nil {
+		return nil, stacktrace.Propagate(err, "Could not upsert USS Availability into repo")
+	}
+	if ussa == nil {
+		return nil, stacktrace.NewError("UpsertUssAvailability returned no USS availability for ID: %s", id)
+	}
+
+	return &restapi.UssAvailabilityStatusResponse{
+		Status:  *ussa.ToRest(),
+		Version: ussa.Version.String(),
+	}, nil
+}

pkg/scd/uss_availability_handler.go

@@ -12,18 +12,8 @@ import (
 	"github.com/interuss/dss/pkg/scd/repos"
 	dssstore "github.com/interuss/dss/pkg/store"
 	"github.com/interuss/stacktrace"
-	"github.com/jackc/pgx/v5"
 )
 
-func GetDefaultAvailabilityResponse(id dssmodels.Manager) *restapi.UssAvailabilityStatusResponse {
-	return &restapi.UssAvailabilityStatusResponse{
-		Status: restapi.UssAvailabilityStatus{
-			Availability: restapi.UssAvailabilityState_Unknown,
-			Uss:          id.String()},
-		Version: "",
-	}
-}
-
 func (a *Server) GetUssAvailability(ctx context.Context, req *restapi.GetUssAvailabilityRequest,
 ) restapi.GetUssAvailabilityResponseSet {
 
@@ -33,30 +23,15 @@ func (a *Server) GetUssAvailability(ctx context.Context, req *restapi.GetUssAvai
 			Message: dsserr.Handle(ctx, stacktrace.NewErrorWithCode(dsserr.BadRequest, "UssId not provided"))}}
 	}
 
-	var response *restapi.UssAvailabilityStatusResponse
-	action := func(ctx context.Context, r repos.Repository) (err error) {
-		// Get USS availability from Store
-		ussa, err := r.GetUssAvailability(ctx, id)
-		if err != nil && err != pgx.ErrNoRows {
-			return stacktrace.Propagate(err, "Could not get USS availability from repo")
-		}
-		if ussa == nil {
-			// Return default availability status "Unknown"
-			response = GetDefaultAvailabilityResponse(id)
-			return nil
-		}
-		response = &restapi.UssAvailabilityStatusResponse{
-			Status:  *ussa.ToRest(),
-			Version: ussa.Version.String(),
-		}
-		return nil
-	}
-
-	_, err := a.Store.Transact(ctx, dssstore.NewFuncOperation(action))
+	response, err := dssstore.TransactWithResult[repos.Repository, *restapi.UssAvailabilityStatusResponse](ctx, a.Store, req)
 	if err != nil {
 		// In case of older DB versions where availability table doesn't exist
 		if strings.Contains(err.Error(), "does not exist") {
-			response = GetDefaultAvailabilityResponse(id)
+			response = &restapi.UssAvailabilityStatusResponse{
+				Status: restapi.UssAvailabilityStatus{
+					Availability: restapi.UssAvailabilityState_Unknown,
+					Uss:          id.String()},
+			}
 		} else {
 			// No need to Propagate this error as this is not a useful stacktrace line
 			return restapi.GetUssAvailabilityResponseSet{Response500: &api.InternalServerErrorBody{
@@ -79,54 +54,21 @@ func (a *Server) SetUssAvailability(ctx context.Context, req *restapi.SetUssAvai
 	}
 
 	// Retrieve USS availability status from request params
-	availability, err := scdmodels.UssAvailabilityStateFromRest(req.Body.Availability)
+	_, err := scdmodels.UssAvailabilityStateFromRest(req.Body.Availability)
 	if err != nil {
 		return restapi.SetUssAvailabilityResponseSet{Response400: &restapi.ErrorResponse{
 			Message: dsserr.Handle(ctx, stacktrace.PropagateWithCode(err, dsserr.BadRequest, "Invalid availability state"))}}
 	}
-	id := dssmodels.ManagerFromString(req.UssId)
-	version := scdmodels.OVN(req.Body.OldVersion)
-	ussareq := &scdmodels.UssAvailabilityStatus{
-		Uss:          id,
-		Availability: availability,
-	}
 
-	var result *restapi.UssAvailabilityStatusResponse
-	action := func(ctx context.Context, r repos.Repository) (err error) {
-		old, err := r.GetUssAvailability(ctx, id)
-		if err != nil && err != pgx.ErrNoRows {
-			return stacktrace.Propagate(err, "Could not get USS availability from repo")
-		}
-		switch {
-		case old == nil && !version.Empty():
-			// The user wants set a new availability status but it already exists.
-			return stacktrace.NewErrorWithCode(dsserr.AlreadyExists, "availability for USS %s already exists", id.String())
-		case old != nil && old.Version != version:
-			// The user wants to update an availability status but the version doesn't match.
-			return stacktrace.Propagate(
-				stacktrace.NewErrorWithCode(dsserr.VersionMismatch, "USS availability version %s is not current", version),
-				"Current version is %s but client specified version %s", old.Version, version)
-		}
-
-		// Upsert the USS availability
-		ussa, err := r.UpsertUssAvailability(ctx, ussareq)
-		if err != nil {
-			return stacktrace.Propagate(err, "Could not upsert USS Availability into repo")
-		}
-		if ussa == nil {
-			return stacktrace.NewError("UpsertUssAvailability returned no USS availability for ID: %s", id)
-		}
-		result = &restapi.UssAvailabilityStatusResponse{
-			Status:  *ussa.ToRest(),
-			Version: ussa.Version.String(),
-		}
-		return nil
-	}
-	_, err = a.Store.Transact(ctx, dssstore.NewFuncOperation(action))
+	result, err := dssstore.TransactWithResult[repos.Repository, *restapi.UssAvailabilityStatusResponse](ctx, a.Store, req)
 	if err != nil {
 		// In case of older DB versions where availability table doesn't exist
 		if strings.Contains(err.Error(), "does not exist") {
-			result = GetDefaultAvailabilityResponse(id)
+			result = &restapi.UssAvailabilityStatusResponse{
+				Status: restapi.UssAvailabilityStatus{
+					Availability: restapi.UssAvailabilityState_Unknown,
+					Uss:          dssmodels.ManagerFromString(req.UssId).String()},
+			}
 		} else {
 			err = stacktrace.Propagate(err, "Could not set USS availability status")
 			errResp := &restapi.ErrorResponse{Message: dsserr.Handle(ctx, err)}