Untitled

mail@pastecode.io avatar
unknown
golang
a year ago
5.6 kB
2
Indexable
Never
package resolvers

// This file will be automatically regenerated based on the schema, any resolver implementations
// will be copied through when generating and any unknown code will be moved to the end.

import (
	"context"
	"fmt"
	"time"

	generated1 "github.com/zeabur/backend/internal/gateway/generated"
	"github.com/zeabur/backend/internal/shared/generated"
	"github.com/zeabur/backend/internal/shared/models"
	pb "github.com/zeabur/backend/internal/shared/proto"
	"github.com/zeabur/backend/internal/shared/utils"
	"go.mongodb.org/mongo-driver/bson/primitive"
	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"google.golang.org/protobuf/types/known/timestamppb"
	"k8s.io/utils/pointer"
)

// Entity is the resolver for the entity field.
func (r *groupedUsageResolver) Entity(ctx context.Context, obj *models.GroupedUsage) (generated.GroupByEntityDetail, error) {
	switch obj.GroupByEntity {
	case models.GroupByEntityProject:
		getProjectReply, err := r.services.Project.GetProject(
			ctx, &pb.GetProjectRequest{
				ProjectID: pointer.String(obj.ID),
			},
		)
		if err != nil {
			r.sugar.Errorw("error getting project for charge", "err", err)
			return nil, fmt.Errorf("error")
		}
		return models.ProjectFromProto(getProjectReply.Project), nil
	case models.GroupByEntityEnvironment:
		getEnvironmentReply, err := r.services.Environment.GetEnvironment(
			ctx, &pb.GetEnvironmentRequest{
				EnvironmentID: obj.ID,
			},
		)
		if err != nil {
			r.sugar.Errorw("error getting environment for charge", "err", err)
			return nil, fmt.Errorf("error")
		}
		return models.EnvironmentFromProto(getEnvironmentReply.Environment), nil
	case models.GroupByEntityService:
		getServiceReply, err := r.services.Project.GetService(
			ctx, &pb.GetServiceRequest{
				ServiceID: pointer.String(obj.ID),
			},
		)
		if err != nil {
			r.sugar.Errorw("error getting service for charge", "err", err)
			return nil, fmt.Errorf("error")
		}
		return models.ServiceFromProto(getServiceReply.Service), nil
	default:
		return nil, fmt.Errorf("unknown group by entity")
	}
}

// Total is the resolver for the total field.
func (r *groupedUsageResolver) Total(ctx context.Context, obj *models.GroupedUsage) ([]float64, error) {
	var results []float64
	protoUsages := obj.ToProto().Usages
	for _, u := range protoUsages {
		results = append(results, u.Usage)
	}
	return results, nil
}

// SingleEnvironmentServiceUsage is the resolver for the singleEnvironmentServiceUsage field.
func (r *queryResolver) SingleEnvironmentServiceUsage(ctx context.Context, serviceID primitive.ObjectID, environmentID primitive.ObjectID, hour string) (*generated.SingleEnvironmentServiceUsage, error) {
	loggedInUserID, isLoggedIn := utils.GetUserId(ctx)
	if !isLoggedIn {
		return nil, fmt.Errorf("not logged in")
	}

	getUserReply, err := r.services.User.GetUser(
		ctx, &pb.GetUserRequest{
			UserID: pointer.String(loggedInUserID.Hex()),
		},
	)
	if err != nil {
		return nil, err
	}

	if getUserReply.User.Id != loggedInUserID.Hex() && !utils.IsAdmin(loggedInUserID) {
		return nil, status.Errorf(codes.PermissionDenied, "not authorized")
	}

	start, err := time.Parse("2006-01-02T15", hour)
	if err != nil {
		return nil, status.Errorf(
			codes.InvalidArgument,
			"invalid month, must be in format YYYY-MM-DDTHH. For example, 2021-01-01T00",
		)
	}

	end := start.Add(time.Hour)

	reply, err := r.services.Metric.QuerySingleEnvironmentServiceUsage(
		ctx, &pb.QuerySingleEnvironmentServiceUsageRequest{
			ServiceID:     serviceID.Hex(),
			EnvironmentID: environmentID.Hex(),
			StartTime:     timestamppb.New(start),
			EndTime:       timestamppb.New(end),
		},
	)
	if err != nil {
		r.sugar.Errorw("failed to get usage", "err", err)
	}

	res := &generated.SingleEnvironmentServiceUsage{
		CPU:    reply.Cpu,
		Memory: int(reply.Memory),
	}

	return res, nil
}

// Usages is the resolver for the usages field.
func (r *queryResolver) Usages(ctx context.Context, from string, to string, userID primitive.ObjectID, groupByEntity *models.GroupByEntity, groupByTime *generated.GroupByTime) (*generated.GroupedResults, error) {
	loggedInUserID, isLoggedIn := utils.GetUserId(ctx)
	if !isLoggedIn {
		return nil, fmt.Errorf("not logged in")
	}

	getUserReply, err := r.services.User.GetUser(
		ctx, &pb.GetUserRequest{
			UserID: pointer.String(loggedInUserID.Hex()),
		},
	)
	if err != nil {
		return nil, err
	}

	if getUserReply.User.Id != loggedInUserID.Hex() && !utils.IsAdmin(loggedInUserID) {
		return nil, status.Errorf(codes.PermissionDenied, "not authorized")
	}

	groupByEntityStr := ""
	groupByTimeStr := ""
	if groupByEntity != nil {
		groupByEntityStr = string(*groupByEntity)
	}
	if groupByTime != nil {
		groupByTimeStr = string(*groupByTime)
	}

	queryChargesReply, err := r.services.Metric.QueryCharges(
		ctx, &pb.QueryChargesRequest{
			UserID:        userID.Hex(),
			StartTime:     from,
			EndTime:       to,
			GroupByEntity: groupByEntityStr,
			GroupByTime:   groupByTimeStr,
		},
	)
	if err != nil {
		return nil, status.Errorf(
			codes.Internal, "failed to query charges",
		)
	}

	r.sugar.Infow("reply", "charge", queryChargesReply)
	charges := make([]*models.GroupedUsage, len(queryChargesReply.Charges))
	for i, charge := range queryChargesReply.Charges {
		charges[i] = models.GroupedUsageFromProto(charge)
	}

	return charges, nil
}

// GroupedUsage returns generated1.GroupedUsageResolver implementation.
func (r *Resolver) GroupedUsage() generated1.GroupedUsageResolver { return &groupedUsageResolver{r} }

type groupedUsageResolver struct{ *Resolver }