feat: add support for get augmented (#2469)

server/resourcemanager/operations.go

@@ -17,6 +17,9 @@ const (
 	OperationUpdate Operation = "update"
 	OperationGet    Operation = "get"
 	OperationDelete Operation = "delete"
+
+	OperationGetAugmented  Operation = "getAugmented"
+	OperationListAugmented Operation = "listAugmented"
 )
 
 var availableOperations = []Operation{
@@ -67,6 +70,14 @@ type Current[T ResourceSpec] interface {
 	Current(context.Context) (T, error)
 }
 
+type GetAugmented[T ResourceSpec] interface {
+	GetAugmented(context.Context, id.ID) (T, error)
+}
+
+type ListAugmented[T ResourceSpec] interface {
+	ListAugmented(_ context.Context, take, skip int, query, sortBy, sortDirection string) ([]T, error)
+}
+
 type resourceHandler[T ResourceSpec] struct {
 	SetID         func(T, id.ID) T
 	List          func(_ context.Context, take, skip int, query, sortBy, sortDirection string) ([]T, error)
@@ -77,6 +88,9 @@ type resourceHandler[T ResourceSpec] struct {
 	Get           func(context.Context, id.ID) (T, error)
 	Delete        func(context.Context, id.ID) error
 	Provision     func(context.Context, T) error
+
+	GetAugmented  func(context.Context, id.ID) (T, error)
+	ListAugmented func(_ context.Context, take, skip int, query, sortBy, sortDirection string) ([]T, error)
 }
 
 func (rh *resourceHandler[T]) bindOperations(enabledOperations []Operation, handler any) error {
@@ -119,6 +133,20 @@ func (rh *resourceHandler[T]) bindOperations(enabledOperations []Operation, hand
 		}
 	}
 
+	if slices.Contains(enabledOperations, OperationGetAugmented) {
+		err := rh.bindGetAugmented(handler)
+		if err != nil {
+			return err
+		}
+	}
+
+	if slices.Contains(enabledOperations, OperationListAugmented) {
+		err := rh.bindListAugmented(handler)
+		if err != nil {
+			return err
+		}
+	}
+
 	err := rh.bindProvisionOperation(handler)
 	if err != nil {
 		return err
@@ -192,3 +220,23 @@ func (rh *resourceHandler[T]) bindProvisionOperation(handler any) error {
 
 	return nil
 }
+
+func (rh *resourceHandler[T]) bindGetAugmented(handler any) error {
+	casted, ok := handler.(GetAugmented[T])
+	if !ok {
+		return fmt.Errorf("handler does not implement interface `GetAugmented[T]`")
+	}
+	rh.GetAugmented = casted.GetAugmented
+
+	return nil
+}
+
+func (rh *resourceHandler[T]) bindListAugmented(handler any) error {
+	casted, ok := handler.(ListAugmented[T])
+	if !ok {
+		return fmt.Errorf("handler does not implement interface `ListAugmented[T]`")
+	}
+	rh.ListAugmented = casted.ListAugmented
+
+	return nil
+}

server/resourcemanager/resource_manager.go

@@ -270,7 +270,12 @@ func (m *manager[T]) list(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	items, err := m.rh.List(
+	listFn := m.rh.List
+	if isRequestForAugmented(r) {
+		listFn = m.rh.ListAugmented
+	}
+
+	items, err := listFn(
 		ctx,
 		take,
 		skip,
@@ -316,6 +321,12 @@ func (m *manager[T]) list(w http.ResponseWriter, r *http.Request) {
 	writeResponse(w, http.StatusOK, string(bytes))
 }
 
+const HeaderAugmented = "X-Tracetest-Augmented"
+
+func isRequestForAugmented(r *http.Request) bool {
+	return r.Header.Get(HeaderAugmented) == "true"
+}
+
 func (m *manager[T]) get(w http.ResponseWriter, r *http.Request) {
 	encoder, err := encoderFromRequest(r)
 	if err != nil {
@@ -327,7 +338,12 @@ func (m *manager[T]) get(w http.ResponseWriter, r *http.Request) {
 	vars := mux.Vars(r)
 	id := id.ID(vars["id"])
 
-	item, err := m.rh.Get(r.Context(), id)
+	getterFn := m.rh.Get
+	if isRequestForAugmented(r) {
+		getterFn = m.rh.GetAugmented
+	}
+
+	item, err := getterFn(r.Context(), id)
 	if err != nil {
 		m.handleResourceHandlerError(w, "getting", err, encoder)
 		return