[usage] introduce getBalance API

components/server/ee/src/billing/billing-service.ts

@@ -40,13 +40,10 @@ export class BillingService {
             };
         }
 
-        const now = new Date();
-        const currentBalance = await this.usageService.listUsage({
+        const getBalanceResponse = await this.usageService.getBalance({
             attributionId: AttributionId.render(attributionId),
-            from: now,
-            to: now,
         });
-        const currentInvoiceCredits = currentBalance.creditBalanceAtEnd | 0;
+        const currentInvoiceCredits = getBalanceResponse.credits;
         if (currentInvoiceCredits >= (costCenter.spendingLimit || 0)) {
             log.info({ userId: user.id }, "Usage limit reached", {
                 attributionId,

components/usage-api/typescript/src/usage/v1/usage.pb.ts

@@ -174,6 +174,14 @@ export interface SetCostCenterRequest {
 export interface SetCostCenterResponse {
 }
 
+export interface GetBalanceRequest {
+  attributionId: string;
+}
+
+export interface GetBalanceResponse {
+  credits: number;
+}
+
 export interface GetCostCenterRequest {
   attributionId: string;
 }
@@ -865,6 +873,100 @@ export const SetCostCenterResponse = {
   },
 };
 
+function createBaseGetBalanceRequest(): GetBalanceRequest {
+  return { attributionId: "" };
+}
+
+export const GetBalanceRequest = {
+  encode(message: GetBalanceRequest, writer: _m0.Writer = _m0.Writer.create()): _m0.Writer {
+    if (message.attributionId !== "") {
+      writer.uint32(10).string(message.attributionId);
+    }
+    return writer;
+  },
+
+  decode(input: _m0.Reader | Uint8Array, length?: number): GetBalanceRequest {
+    const reader = input instanceof _m0.Reader ? input : new _m0.Reader(input);
+    let end = length === undefined ? reader.len : reader.pos + length;
+    const message = createBaseGetBalanceRequest();
+    while (reader.pos < end) {
+      const tag = reader.uint32();
+      switch (tag >>> 3) {
+        case 1:
+          message.attributionId = reader.string();
+          break;
+        default:
+          reader.skipType(tag & 7);
+          break;
+      }
+    }
+    return message;
+  },
+
+  fromJSON(object: any): GetBalanceRequest {
+    return { attributionId: isSet(object.attributionId) ? String(object.attributionId) : "" };
+  },
+
+  toJSON(message: GetBalanceRequest): unknown {
+    const obj: any = {};
+    message.attributionId !== undefined && (obj.attributionId = message.attributionId);
+    return obj;
+  },
+
+  fromPartial(object: DeepPartial<GetBalanceRequest>): GetBalanceRequest {
+    const message = createBaseGetBalanceRequest();
+    message.attributionId = object.attributionId ?? "";
+    return message;
+  },
+};
+
+function createBaseGetBalanceResponse(): GetBalanceResponse {
+  return { credits: 0 };
+}
+
+export const GetBalanceResponse = {
+  encode(message: GetBalanceResponse, writer: _m0.Writer = _m0.Writer.create()): _m0.Writer {
+    if (message.credits !== 0) {
+      writer.uint32(33).double(message.credits);
+    }
+    return writer;
+  },
+
+  decode(input: _m0.Reader | Uint8Array, length?: number): GetBalanceResponse {
+    const reader = input instanceof _m0.Reader ? input : new _m0.Reader(input);
+    let end = length === undefined ? reader.len : reader.pos + length;
+    const message = createBaseGetBalanceResponse();
+    while (reader.pos < end) {
+      const tag = reader.uint32();
+      switch (tag >>> 3) {
+        case 4:
+          message.credits = reader.double();
+          break;
+        default:
+          reader.skipType(tag & 7);
+          break;
+      }
+    }
+    return message;
+  },
+
+  fromJSON(object: any): GetBalanceResponse {
+    return { credits: isSet(object.credits) ? Number(object.credits) : 0 };
+  },
+
+  toJSON(message: GetBalanceResponse): unknown {
+    const obj: any = {};
+    message.credits !== undefined && (obj.credits = message.credits);
+    return obj;
+  },
+
+  fromPartial(object: DeepPartial<GetBalanceResponse>): GetBalanceResponse {
+    const message = createBaseGetBalanceResponse();
+    message.credits = object.credits ?? 0;
+    return message;
+  },
+};
+
 function createBaseGetCostCenterRequest(): GetCostCenterRequest {
   return { attributionId: "" };
 }
@@ -1087,6 +1189,15 @@ export const UsageServiceDefinition = {
       responseStream: false,
       options: {},
     },
+    /** GetBalance returns the current credits balance for the given attributionId */
+    getBalance: {
+      name: "GetBalance",
+      requestType: GetBalanceRequest,
+      requestStream: false,
+      responseType: GetBalanceResponse,
+      responseStream: false,
+      options: {},
+    },
   },
 } as const;
 
@@ -1108,6 +1219,11 @@ export interface UsageServiceServiceImplementation<CallContextExt = {}> {
   ): Promise<DeepPartial<ReconcileUsageResponse>>;
   /** ListUsage retrieves all usage for the specified attributionId and theb given time range */
   listUsage(request: ListUsageRequest, context: CallContext & CallContextExt): Promise<DeepPartial<ListUsageResponse>>;
+  /** GetBalance returns the current credits balance for the given attributionId */
+  getBalance(
+    request: GetBalanceRequest,
+    context: CallContext & CallContextExt,
+  ): Promise<DeepPartial<GetBalanceResponse>>;
 }
 
 export interface UsageServiceClient<CallOptionsExt = {}> {
@@ -1128,6 +1244,11 @@ export interface UsageServiceClient<CallOptionsExt = {}> {
   ): Promise<ReconcileUsageResponse>;
   /** ListUsage retrieves all usage for the specified attributionId and theb given time range */
   listUsage(request: DeepPartial<ListUsageRequest>, options?: CallOptions & CallOptionsExt): Promise<ListUsageResponse>;
+  /** GetBalance returns the current credits balance for the given attributionId */
+  getBalance(
+    request: DeepPartial<GetBalanceRequest>,
+    options?: CallOptions & CallOptionsExt,
+  ): Promise<GetBalanceResponse>;
 }
 
 export interface DataLoaderOptions {

components/usage-api/usage/v1/usage.proto

@@ -19,6 +19,9 @@ service UsageService {
 
     // ListUsage retrieves all usage for the specified attributionId and theb given time range
     rpc ListUsage(ListUsageRequest) returns (ListUsageResponse) {}
+
+    // GetBalance returns the current credits balance for the given attributionId
+    rpc GetBalance(GetBalanceRequest) returns (GetBalanceResponse) {}
 }
 
 message ReconcileUsageRequest {
@@ -98,6 +101,14 @@ message SetCostCenterRequest {
 message SetCostCenterResponse {
 }
 
+message GetBalanceRequest {
+    string attribution_id = 1;
+}
+
+message GetBalanceResponse {
+    double credits = 4;
+}
+
 message GetCostCenterRequest {
     string attribution_id = 1;
 }

components/usage/pkg/apiv1/usage.go

@@ -156,6 +156,20 @@ func (s *UsageService) ListUsage(ctx context.Context, in *v1.ListUsageRequest) (
 	}, nil
 }
 
+func (s *UsageService) GetBalance(ctx context.Context, in *v1.GetBalanceRequest) (*v1.GetBalanceResponse, error) {
+	attrId, err := db.ParseAttributionID(in.AttributionId)
+	if err != nil {
+		return nil, err
+	}
+	credits, err := db.GetBalance(ctx, s.conn, attrId)
+	if err != nil {
+		return nil, err
+	}
+	return &v1.GetBalanceResponse{
+		Credits: credits.ToCredits(),
+	}, nil
+}
+
 func (s *UsageService) GetCostCenter(ctx context.Context, in *v1.GetCostCenterRequest) (*v1.GetCostCenterResponse, error) {
 	if in.AttributionId == "" {
 		return nil, status.Errorf(codes.InvalidArgument, "Empty attributionId")

components/usage/pkg/db/usage.go

@@ -207,6 +207,30 @@ type Balance struct {
 	CreditCents   CreditCents   `gorm:"column:creditCents;type:bigint;" json:"creditCents"`
 }
 
+func GetBalance(ctx context.Context, conn *gorm.DB, attributionId AttributionID) (CreditCents, error) {
+	rows, err := conn.WithContext(ctx).
+		Model(&Usage{}).
+		Select("sum(creditCents) as balance").
+		Where("attributionId = ?", string(attributionId)).
+		Group("attributionId").
+		Rows()
+	if err != nil {
+		return 0, fmt.Errorf("failed to get rows for list balance query: %w", err)
+	}
+	defer rows.Close()
+
+	if !rows.Next() {
+		return 0, nil
+	}
+
+	var balance CreditCents
+	err = conn.ScanRows(rows, &balance)
+	if err != nil {
+		return 0, fmt.Errorf("failed to scan row: %w", err)
+	}
+	return balance, nil
+}
+
 func ListBalance(ctx context.Context, conn *gorm.DB) ([]Balance, error) {
 	var balances []Balance
 	rows, err := conn.WithContext(ctx).

components/usage/pkg/db/usage_test.go

@@ -322,3 +322,37 @@ func TestListBalance(t *testing.T) {
 		CreditCents:   -50,
 	})
 }
+
+func TestGetBalance(t *testing.T) {
+	teamAttributionID := db.NewTeamAttributionID(uuid.New().String())
+	userAttributionID := db.NewUserAttributionID(uuid.New().String())
+
+	conn := dbtest.ConnectForTests(t)
+	dbtest.CreateUsageRecords(t, conn,
+		dbtest.NewUsage(t, db.Usage{
+			AttributionID: teamAttributionID,
+			CreditCents:   100,
+		}),
+		dbtest.NewUsage(t, db.Usage{
+			AttributionID: teamAttributionID,
+			CreditCents:   900,
+		}),
+		dbtest.NewUsage(t, db.Usage{
+			AttributionID: userAttributionID,
+			CreditCents:   450,
+		}),
+		dbtest.NewUsage(t, db.Usage{
+			AttributionID: userAttributionID,
+			CreditCents:   -500,
+			Kind:          db.InvoiceUsageKind,
+		}),
+	)
+
+	teamBalance, err := db.GetBalance(context.Background(), conn, teamAttributionID)
+	require.NoError(t, err)
+	require.EqualValues(t, 1000, int(teamBalance))
+
+	userBalance, err := db.GetBalance(context.Background(), conn, userAttributionID)
+	require.NoError(t, err)
+	require.EqualValues(t, -50, int(userBalance))
+}