refactor(billing): Billing simplification entitlement usage

Makefile

@@ -4,7 +4,7 @@ TAG := $(shell git rev-list --tags --max-count=1)
 VERSION := $(shell git describe --tags ${TAG})
 .PHONY: build check fmt lint test test-race vet test-cover-html help install proto ui compose-up-dev
 .DEFAULT_GOAL := build
-PROTON_COMMIT := "80fc5ba1e538e38d5ca190386af1e69ee64584ee"
+PROTON_COMMIT := "4144445eb0f9cbd1a801a3d0aa5cfce4cc0ea551"
 
 ui:
 	@echo " > generating ui build"

internal/api/v1beta1connect/billing_check.go

@@ -2,6 +2,7 @@ package v1beta1connect
 
 import (
 	"context"
+	"errors"
 
 	"connectrpc.com/connect"
 	"github.com/raystack/frontier/billing/customer"
@@ -12,10 +13,25 @@ import (
 func (h *ConnectHandler) CheckFeatureEntitlement(ctx context.Context, request *connect.Request[frontierv1beta1.CheckFeatureEntitlementRequest]) (*connect.Response[frontierv1beta1.CheckFeatureEntitlementResponse], error) {
 	errorLogger := NewErrorLogger()
 
-	checkStatus, err := h.entitlementService.Check(ctx, request.Msg.GetBillingId(), request.Msg.GetFeature())
+	// Always infer billing_id from org_id
+	cust, err := h.customerService.GetByOrgID(ctx, request.Msg.GetOrgId())
+	if err != nil {
+		if errors.Is(err, customer.ErrNotFound) {
+			return connect.NewResponse(&frontierv1beta1.CheckFeatureEntitlementResponse{}), nil
+		}
+		if errors.Is(err, customer.ErrInvalidUUID) || errors.Is(err, customer.ErrInvalidID) {
+			return nil, connect.NewError(connect.CodeInvalidArgument, err)
+		}
+		errorLogger.LogServiceError(ctx, request, "CheckFeatureEntitlement.GetByOrgID", err,
+			zap.String("org_id", request.Msg.GetOrgId()))
+		return nil, connect.NewError(connect.CodeInternal, ErrInternalServerError)
+	}
+
+	checkStatus, err := h.entitlementService.Check(ctx, cust.ID, request.Msg.GetFeature())
 	if err != nil {
 		errorLogger.LogServiceError(ctx, request, "CheckFeatureEntitlement", err,
-			zap.String("billing_id", request.Msg.GetBillingId()),
+			zap.String("billing_id", cust.ID),
+			zap.String("org_id", request.Msg.GetOrgId()),
 			zap.String("feature", request.Msg.GetFeature()))
 		return nil, connect.NewError(connect.CodeInternal, ErrInternalServerError)
 	}

internal/api/v1beta1connect/billing_check_test.go

@@ -15,19 +15,23 @@ import (
 
 func TestConnectHandler_CheckFeatureEntitlement(t *testing.T) {
 	tests := []struct {
-		name    string
-		setup   func(es *mocks.EntitlementService)
-		request *connect.Request[frontierv1beta1.CheckFeatureEntitlementRequest]
-		want    *connect.Response[frontierv1beta1.CheckFeatureEntitlementResponse]
-		wantErr error
-		errCode connect.Code
+		name          string
+		customerSetup func(cs *mocks.CustomerService)
+		setup         func(es *mocks.EntitlementService)
+		request       *connect.Request[frontierv1beta1.CheckFeatureEntitlementRequest]
+		want          *connect.Response[frontierv1beta1.CheckFeatureEntitlementResponse]
+		wantErr       error
+		errCode       connect.Code
 	}{
 		{
 			name: "should return internal server error when entitlement service returns error",
 			request: connect.NewRequest(&frontierv1beta1.CheckFeatureEntitlementRequest{
-				BillingId: "billing-123",
-				Feature:   "feature-abc",
+				OrgId:   "org-123",
+				Feature: "feature-abc",
 			}),
+			customerSetup: func(cs *mocks.CustomerService) {
+				cs.EXPECT().GetByOrgID(mock.Anything, "org-123").Return(customer.Customer{ID: "billing-123"}, nil)
+			},
 			setup: func(es *mocks.EntitlementService) {
 				es.EXPECT().Check(mock.Anything, "billing-123", "feature-abc").Return(false, errors.New("service error"))
 			},
@@ -37,13 +41,16 @@ func TestConnectHandler_CheckFeatureEntitlement(t *testing.T) {
 		},
 		{
 			name: "should return false when feature is not entitled",
+			request: connect.NewRequest(&frontierv1beta1.CheckFeatureEntitlementRequest{
+				OrgId:   "org-123",
+				Feature: "feature-abc",
+			}),
+			customerSetup: func(cs *mocks.CustomerService) {
+				cs.EXPECT().GetByOrgID(mock.Anything, "org-123").Return(customer.Customer{ID: "billing-123"}, nil)
+			},
 			setup: func(es *mocks.EntitlementService) {
 				es.EXPECT().Check(mock.Anything, "billing-123", "feature-abc").Return(false, nil)
 			},
-			request: connect.NewRequest(&frontierv1beta1.CheckFeatureEntitlementRequest{
-				BillingId: "billing-123",
-				Feature:   "feature-abc",
-			}),
 			want: connect.NewResponse(&frontierv1beta1.CheckFeatureEntitlementResponse{
 				Status: false,
 			}),
@@ -52,57 +59,63 @@ func TestConnectHandler_CheckFeatureEntitlement(t *testing.T) {
 		},
 		{
 			name: "should return true when feature is entitled",
+			request: connect.NewRequest(&frontierv1beta1.CheckFeatureEntitlementRequest{
+				OrgId:   "org-123",
+				Feature: "feature-abc",
+			}),
+			customerSetup: func(cs *mocks.CustomerService) {
+				cs.EXPECT().GetByOrgID(mock.Anything, "org-123").Return(customer.Customer{ID: "billing-123"}, nil)
+			},
 			setup: func(es *mocks.EntitlementService) {
 				es.EXPECT().Check(mock.Anything, "billing-123", "feature-abc").Return(true, nil)
 			},
-			request: connect.NewRequest(&frontierv1beta1.CheckFeatureEntitlementRequest{
-				BillingId: "billing-123",
-				Feature:   "feature-abc",
-			}),
 			want: connect.NewResponse(&frontierv1beta1.CheckFeatureEntitlementResponse{
 				Status: true,
 			}),
 			wantErr: nil,
 			errCode: connect.Code(0),
 		},
 		{
-			name: "should handle empty billing id",
-			setup: func(es *mocks.EntitlementService) {
-				es.EXPECT().Check(mock.Anything, "", "feature-abc").Return(false, nil)
-			},
+			name: "should return empty response when billing account not found",
 			request: connect.NewRequest(&frontierv1beta1.CheckFeatureEntitlementRequest{
-				BillingId: "",
-				Feature:   "feature-abc",
-			}),
-			want: connect.NewResponse(&frontierv1beta1.CheckFeatureEntitlementResponse{
-				Status: false,
+				OrgId:   "org-123",
+				Feature: "feature-abc",
 			}),
+			customerSetup: func(cs *mocks.CustomerService) {
+				cs.EXPECT().GetByOrgID(mock.Anything, "org-123").Return(customer.Customer{}, customer.ErrNotFound)
+			},
+			setup:   func(es *mocks.EntitlementService) {},
+			want:    connect.NewResponse(&frontierv1beta1.CheckFeatureEntitlementResponse{}),
 			wantErr: nil,
 			errCode: connect.Code(0),
 		},
 		{
-			name: "should handle empty feature",
-			setup: func(es *mocks.EntitlementService) {
-				es.EXPECT().Check(mock.Anything, "billing-123", "").Return(false, nil)
-			},
+			name: "should return invalid argument when org_id is invalid",
 			request: connect.NewRequest(&frontierv1beta1.CheckFeatureEntitlementRequest{
-				BillingId: "billing-123",
-				Feature:   "",
+				OrgId:   "",
+				Feature: "feature-abc",
 			}),
-			want: connect.NewResponse(&frontierv1beta1.CheckFeatureEntitlementResponse{
-				Status: false,
-			}),
-			wantErr: nil,
-			errCode: connect.Code(0),
+			customerSetup: func(cs *mocks.CustomerService) {
+				cs.EXPECT().GetByOrgID(mock.Anything, "").Return(customer.Customer{}, customer.ErrInvalidUUID)
+			},
+			setup:   func(es *mocks.EntitlementService) {},
+			want:    nil,
+			wantErr: customer.ErrInvalidUUID,
+			errCode: connect.CodeInvalidArgument,
 		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			mockCustomerSvc := new(mocks.CustomerService)
 			mockEntitlementSvc := new(mocks.EntitlementService)
+			if tt.customerSetup != nil {
+				tt.customerSetup(mockCustomerSvc)
+			}
 			if tt.setup != nil {
 				tt.setup(mockEntitlementSvc)
 			}
 			h := &ConnectHandler{
+				customerService:    mockCustomerSvc,
 				entitlementService: mockEntitlementSvc,
 			}
 			got, err := h.CheckFeatureEntitlement(context.Background(), tt.request)

internal/api/v1beta1connect/billing_usage.go

@@ -20,6 +20,20 @@ import (
 func (h *ConnectHandler) CreateBillingUsage(ctx context.Context, request *connect.Request[frontierv1beta1.CreateBillingUsageRequest]) (*connect.Response[frontierv1beta1.CreateBillingUsageResponse], error) {
 	errorLogger := NewErrorLogger()
 
+	// Always infer billing_id from org_id
+	cust, err := h.customerService.GetByOrgID(ctx, request.Msg.GetOrgId())
+	if err != nil {
+		if errors.Is(err, customer.ErrNotFound) {
+			return nil, connect.NewError(connect.CodeNotFound, err)
+		}
+		if errors.Is(err, customer.ErrInvalidUUID) || errors.Is(err, customer.ErrInvalidID) {
+			return nil, connect.NewError(connect.CodeInvalidArgument, err)
+		}
+		errorLogger.LogServiceError(ctx, request, "CreateBillingUsage.GetByOrgID", err,
+			zap.String("org_id", request.Msg.GetOrgId()))
+		return nil, connect.NewError(connect.CodeInternal, ErrInternalServerError)
+	}
+
 	createRequests := make([]usage.Usage, 0, len(request.Msg.GetUsages()))
 	for _, v := range request.Msg.GetUsages() {
 		usageType := usage.CreditType
@@ -29,7 +43,7 @@ func (h *ConnectHandler) CreateBillingUsage(ctx context.Context, request *connec
 
 		createRequests = append(createRequests, usage.Usage{
 			ID:          v.GetId(),
-			CustomerID:  request.Msg.GetBillingId(),
+			CustomerID:  cust.ID,
 			Type:        usageType,
 			Amount:      v.GetAmount(),
 			Source:      strings.ToLower(v.GetSource()), // source in lower case looks nicer
@@ -47,7 +61,8 @@ func (h *ConnectHandler) CreateBillingUsage(ctx context.Context, request *connec
 			return nil, connect.NewError(connect.CodeAlreadyExists, ErrAlreadyApplied)
 		}
 		errorLogger.LogServiceError(ctx, request, "CreateBillingUsage.Report", err,
-			zap.String("billing_id", request.Msg.GetBillingId()),
+			zap.String("billing_id", cust.ID),
+			zap.String("org_id", request.Msg.GetOrgId()),
 			zap.Int("usage_count", len(createRequests)))
 		return nil, connect.NewError(connect.CodeInternal, ErrInternalServerError)
 	}
@@ -185,7 +200,21 @@ func transformTransactionToPB(t credit.Transaction) (*frontierv1beta1.BillingTra
 func (h *ConnectHandler) RevertBillingUsage(ctx context.Context, request *connect.Request[frontierv1beta1.RevertBillingUsageRequest]) (*connect.Response[frontierv1beta1.RevertBillingUsageResponse], error) {
 	errorLogger := NewErrorLogger()
 
-	if err := h.usageService.Revert(ctx, request.Msg.GetBillingId(),
+	// Always infer billing_id from org_id
+	cust, err := h.customerService.GetByOrgID(ctx, request.Msg.GetOrgId())
+	if err != nil {
+		if errors.Is(err, customer.ErrNotFound) {
+			return nil, connect.NewError(connect.CodeNotFound, err)
+		}
+		if errors.Is(err, customer.ErrInvalidUUID) || errors.Is(err, customer.ErrInvalidID) {
+			return nil, connect.NewError(connect.CodeInvalidArgument, err)
+		}
+		errorLogger.LogServiceError(ctx, request, "RevertBillingUsage.GetByOrgID", err,
+			zap.String("org_id", request.Msg.GetOrgId()))
+		return nil, connect.NewError(connect.CodeInternal, ErrInternalServerError)
+	}
+
+	if err := h.usageService.Revert(ctx, cust.ID,
 		request.Msg.GetUsageId(), request.Msg.GetAmount()); err != nil {
 		if errors.Is(err, usage.ErrRevertAmountExceeds) {
 			return nil, connect.NewError(connect.CodeInvalidArgument, err)
@@ -199,7 +228,8 @@ func (h *ConnectHandler) RevertBillingUsage(ctx context.Context, request *connec
 			return nil, connect.NewError(connect.CodeInvalidArgument, err)
 		}
 		errorLogger.LogServiceError(ctx, request, "RevertBillingUsage.Revert", err,
-			zap.String("billing_id", request.Msg.GetBillingId()),
+			zap.String("billing_id", cust.ID),
+			zap.String("org_id", request.Msg.GetOrgId()),
 			zap.String("usage_id", request.Msg.GetUsageId()),
 			zap.Int64("amount", request.Msg.GetAmount()))
 		return nil, connect.NewError(connect.CodeInternal, ErrInternalServerError)