allow setting org level billing related banners

admin/database/database.go

@@ -1316,6 +1316,7 @@ const (
 	BillingIssueTypeNeverSubscribed                        = 7
 	BillingIssueTypeOnCreditTrial                          = 8
 	BillingIssueTypeTrialCreditsDepleted                   = 9
+	BillingIssueTypeMessage                                = 10
 )
 
 type BillingIssueLevel int
@@ -1390,9 +1391,16 @@ type BillingIssueMetadataTrialCreditsDepleted struct {
 	DepletedOn time.Time `json:"depleted_on"`
 }
 
+type BillingIssueMetadataMessage struct {
+	Message string `json:"message"`
+}
+
 type UpsertBillingIssueOptions struct {
-	OrgID     string           `validate:"required"`
-	Type      BillingIssueType `validate:"required"`
+	OrgID string           `validate:"required"`
+	Type  BillingIssueType `validate:"required"`
+	// Level is optional; if unspecified, it is derived from Type. It is used for types whose
+	// severity is not implied by the type itself (e.g. BillingIssueTypeMessage).
+	Level     BillingIssueLevel `exhaustruct:"optional"`
 	Metadata  BillingIssueMetadata
 	EventTime time.Time `validate:"required"`
 }

admin/database/postgres/postgres.go

@@ -3122,15 +3122,18 @@ func (c *connection) UpsertBillingIssue(ctx context.Context, opts *database.Upse
 	temp := &billingIssueDTO{
 		OrgID:     opts.OrgID,
 		Type:      opts.Type,
+		Level:     opts.Level,
 		Metadata:  metadata,
 		EventTime: opts.EventTime,
 	}
 
-	temp.Level = temp.getBillingIssueLevel()
+	if temp.Level == database.BillingIssueLevelUnspecified {
+		temp.Level = temp.getBillingIssueLevel()
+	}
 
 	res := &billingIssueDTO{}
 	err = c.getDB(ctx).QueryRowxContext(ctx, `INSERT INTO billing_issues (org_id, type, level, metadata, event_time) VALUES ($1, $2, $3, $4, $5)
-		ON CONFLICT (org_id, type) DO UPDATE SET metadata = $4, event_time = $5 RETURNING *`, temp.OrgID, temp.Type, temp.Level, temp.Metadata, temp.EventTime).StructScan(res)
+		ON CONFLICT (org_id, type) DO UPDATE SET level = $3, metadata = $4, event_time = $5 RETURNING *`, temp.OrgID, temp.Type, temp.Level, temp.Metadata, temp.EventTime).StructScan(res)
 	if err != nil {
 		return nil, parseErr("billing issue", err)
 	}
@@ -3753,6 +3756,8 @@ func (b *billingIssueDTO) AsModel() *database.BillingIssue {
 		metadata = &database.BillingIssueMetadataOnCreditTrial{}
 	case database.BillingIssueTypeTrialCreditsDepleted:
 		metadata = &database.BillingIssueMetadataTrialCreditsDepleted{}
+	case database.BillingIssueTypeMessage:
+		metadata = &database.BillingIssueMetadataMessage{}
 	default:
 	}
 	if err := json.Unmarshal(b.Metadata, &metadata); err != nil {

admin/server/billing.go

@@ -991,6 +991,63 @@ func (s *Server) SudoDeleteOrganizationBillingIssue(ctx context.Context, req *ad
 	return &adminv1.SudoDeleteOrganizationBillingIssueResponse{}, nil
 }
 
+func (s *Server) SudoUpdateOrganizationBillingMessage(ctx context.Context, req *adminv1.SudoUpdateOrganizationBillingMessageRequest) (*adminv1.SudoUpdateOrganizationBillingMessageResponse, error) {
+	observability.AddRequestAttributes(ctx, attribute.String("args.org", req.Org), attribute.String("args.level", req.Level.String()))
+
+	claims := auth.GetClaims(ctx)
+	if !claims.Superuser(ctx) {
+		return nil, status.Error(codes.PermissionDenied, "only superusers can set the billing message")
+	}
+
+	if req.Message == "" {
+		return nil, status.Error(codes.InvalidArgument, "message is required")
+	}
+
+	level, err := dtoBillingIssueLevelToDB(req.Level)
+	if err != nil {
+		return nil, err
+	}
+
+	org, err := s.admin.DB.FindOrganizationByName(ctx, req.Org)
+	if err != nil {
+		return nil, err
+	}
+
+	_, err = s.admin.DB.UpsertBillingIssue(ctx, &database.UpsertBillingIssueOptions{
+		OrgID:     org.ID,
+		Type:      database.BillingIssueTypeMessage,
+		Level:     level,
+		Metadata:  &database.BillingIssueMetadataMessage{Message: req.Message},
+		EventTime: time.Now(),
+	})
+	if err != nil {
+		return nil, err
+	}
+
+	return &adminv1.SudoUpdateOrganizationBillingMessageResponse{}, nil
+}
+
+func (s *Server) SudoDeleteOrganizationBillingMessage(ctx context.Context, req *adminv1.SudoDeleteOrganizationBillingMessageRequest) (*adminv1.SudoDeleteOrganizationBillingMessageResponse, error) {
+	observability.AddRequestAttributes(ctx, attribute.String("args.org", req.Org))
+
+	claims := auth.GetClaims(ctx)
+	if !claims.Superuser(ctx) {
+		return nil, status.Error(codes.PermissionDenied, "only superusers can delete the billing message")
+	}
+
+	org, err := s.admin.DB.FindOrganizationByName(ctx, req.Org)
+	if err != nil {
+		return nil, err
+	}
+
+	err = s.admin.DB.DeleteBillingIssueByTypeForOrg(ctx, org.ID, database.BillingIssueTypeMessage)
+	if err != nil {
+		return nil, err
+	}
+
+	return &adminv1.SudoDeleteOrganizationBillingMessageResponse{}, nil
+}
+
 func (s *Server) updateQuotasAndHandleBillingIssues(ctx context.Context, org *database.Organization, sub *billing.Subscription) (*database.Organization, error) {
 	org, err := s.admin.DB.UpdateOrganization(ctx, org.ID, &database.UpdateOrganizationOptions{
 		Name:                                org.Name,
@@ -1168,6 +1225,8 @@ func billingIssueTypeToDTO(t database.BillingIssueType) adminv1.BillingIssueType
 		return adminv1.BillingIssueType_BILLING_ISSUE_TYPE_ON_CREDIT_TRIAL
 	case database.BillingIssueTypeTrialCreditsDepleted:
 		return adminv1.BillingIssueType_BILLING_ISSUE_TYPE_TRIAL_CREDITS_DEPLETED
+	case database.BillingIssueTypeMessage:
+		return adminv1.BillingIssueType_BILLING_ISSUE_TYPE_MESSAGE
 	default:
 		return adminv1.BillingIssueType_BILLING_ISSUE_TYPE_UNSPECIFIED
 	}
@@ -1184,6 +1243,17 @@ func billingIssueLevelToDTO(l database.BillingIssueLevel) adminv1.BillingIssueLe
 	}
 }
 
+func dtoBillingIssueLevelToDB(l adminv1.BillingIssueLevel) (database.BillingIssueLevel, error) {
+	switch l {
+	case adminv1.BillingIssueLevel_BILLING_ISSUE_LEVEL_WARNING:
+		return database.BillingIssueLevelWarning, nil
+	case adminv1.BillingIssueLevel_BILLING_ISSUE_LEVEL_ERROR:
+		return database.BillingIssueLevelError, nil
+	default:
+		return database.BillingIssueLevelUnspecified, status.Error(codes.InvalidArgument, "invalid billing issue level")
+	}
+}
+
 func dtoBillingIssueTypeToDB(t adminv1.BillingIssueType) (database.BillingIssueType, error) {
 	switch t {
 	case adminv1.BillingIssueType_BILLING_ISSUE_TYPE_ON_TRIAL:
@@ -1204,6 +1274,8 @@ func dtoBillingIssueTypeToDB(t adminv1.BillingIssueType) (database.BillingIssueT
 		return database.BillingIssueTypeOnCreditTrial, nil
 	case adminv1.BillingIssueType_BILLING_ISSUE_TYPE_TRIAL_CREDITS_DEPLETED:
 		return database.BillingIssueTypeTrialCreditsDepleted, nil
+	case adminv1.BillingIssueType_BILLING_ISSUE_TYPE_MESSAGE:
+		return database.BillingIssueTypeMessage, nil
 	default:
 		return database.BillingIssueTypeUnspecified, status.Error(codes.InvalidArgument, "invalid billing error type")
 	}
@@ -1318,6 +1390,14 @@ func billingIssueMetadataToDTO(t database.BillingIssueType, m database.BillingIs
 				},
 			},
 		}
+	case database.BillingIssueTypeMessage:
+		return &adminv1.BillingIssueMetadata{
+			Metadata: &adminv1.BillingIssueMetadata_Message{
+				Message: &adminv1.BillingIssueMetadataMessage{
+					Message: m.(*database.BillingIssueMetadataMessage).Message,
+				},
+			},
+		}
 	default:
 		return &adminv1.BillingIssueMetadata{}
 	}

cli/cmd/sudo/billing/billing.go

@@ -13,6 +13,8 @@ func BillingCmd(ch *cmdutil.Helper) *cobra.Command {
 
 	billingCmd.AddCommand(SetCmd(ch))
 	billingCmd.AddCommand(DeleteIssueCmd(ch))
+	billingCmd.AddCommand(SetMessageCmd(ch))
+	billingCmd.AddCommand(DeleteMessageCmd(ch))
 	billingCmd.AddCommand(GrantTrialCreditsCmd(ch))
 	billingCmd.AddCommand(RepairCmd(ch))
 	billingCmd.AddCommand(SetupCmd(ch))

cli/cmd/sudo/billing/delete-message.go

@@ -0,0 +1,44 @@
+package billing
+
+import (
+	"fmt"
+
+	"github.com/rilldata/rill/cli/pkg/cmdutil"
+	adminv1 "github.com/rilldata/rill/proto/gen/rill/admin/v1"
+	"github.com/spf13/cobra"
+)
+
+func DeleteMessageCmd(ch *cmdutil.Helper) *cobra.Command {
+	var org string
+	cmd := &cobra.Command{
+		Use:   "delete-message",
+		Short: "Remove the message banner for an organization",
+		Args:  cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+
+			client, err := ch.Client()
+			if err != nil {
+				return err
+			}
+
+			if org == "" {
+				return fmt.Errorf("please set --org")
+			}
+
+			_, err = client.SudoDeleteOrganizationBillingMessage(ctx, &adminv1.SudoDeleteOrganizationBillingMessageRequest{
+				Org: org,
+			})
+			if err != nil {
+				return err
+			}
+
+			ch.PrintfSuccess("Message banner removed for organization %q\n", org)
+
+			return nil
+		},
+	}
+
+	cmd.Flags().StringVar(&org, "org", "", "Organization Name")
+	return cmd
+}

cli/cmd/sudo/billing/set-message.go

@@ -0,0 +1,73 @@
+package billing
+
+import (
+	"fmt"
+
+	"github.com/rilldata/rill/cli/pkg/cmdutil"
+	adminv1 "github.com/rilldata/rill/proto/gen/rill/admin/v1"
+	"github.com/spf13/cobra"
+)
+
+func SetMessageCmd(ch *cmdutil.Helper) *cobra.Command {
+	var org, level, message string
+	levels := []string{"warning", "error"}
+	cmd := &cobra.Command{
+		Use:   "set-message",
+		Short: "Set (overriding any existing) a message banner for an organization",
+		Args:  cobra.NoArgs,
+		RunE: func(cmd *cobra.Command, args []string) error {
+			ctx := cmd.Context()
+
+			client, err := ch.Client()
+			if err != nil {
+				return err
+			}
+
+			if org == "" {
+				return fmt.Errorf("please set --org")
+			}
+
+			if message == "" {
+				return fmt.Errorf("please set --message")
+			}
+
+			if level == "" {
+				if !ch.Interactive {
+					return fmt.Errorf("--level flag is required in non-interactive mode")
+				}
+				level, err = cmdutil.SelectPrompt("Select message level", levels, "warning")
+				if err != nil {
+					return err
+				}
+			}
+
+			var l adminv1.BillingIssueLevel
+			switch level {
+			case "warning":
+				l = adminv1.BillingIssueLevel_BILLING_ISSUE_LEVEL_WARNING
+			case "error":
+				l = adminv1.BillingIssueLevel_BILLING_ISSUE_LEVEL_ERROR
+			default:
+				return fmt.Errorf("invalid level %q, must be one of: warning, error", level)
+			}
+
+			_, err = client.SudoUpdateOrganizationBillingMessage(ctx, &adminv1.SudoUpdateOrganizationBillingMessageRequest{
+				Org:     org,
+				Level:   l,
+				Message: message,
+			})
+			if err != nil {
+				return err
+			}
+
+			ch.PrintfSuccess("Message banner set for organization %q\n", org)
+
+			return nil
+		},
+	}
+
+	cmd.Flags().StringVar(&org, "org", "", "Organization Name")
+	cmd.Flags().StringVar(&level, "level", "", "Message level (warning or error)")
+	cmd.Flags().StringVar(&message, "message", "", "Message to display in the banner")
+	return cmd
+}

proto/gen/rill/admin/v1/admin.swagger.yaml

@@ -4023,6 +4023,52 @@ paths:
             - BILLING_ISSUE_TYPE_NEVER_SUBSCRIBED
             - BILLING_ISSUE_TYPE_ON_CREDIT_TRIAL
             - BILLING_ISSUE_TYPE_TRIAL_CREDITS_DEPLETED
+            - BILLING_ISSUE_TYPE_MESSAGE
+  /v1/superuser/organizations/{org}/billing/message:
+    delete:
+      summary: SudoDeleteOrganizationBillingMessage removes the custom message banner for the organization
+      operationId: AdminService_SudoDeleteOrganizationBillingMessage
+      responses:
+        "200":
+          description: A successful response.
+          schema:
+            $ref: '#/definitions/v1SudoDeleteOrganizationBillingMessageResponse'
+        default:
+          description: An unexpected error response.
+          schema:
+            $ref: '#/definitions/rpcStatus'
+      parameters:
+        - name: org
+          in: path
+          required: true
+          type: string
+    post:
+      summary: SudoUpdateOrganizationBillingMessage sets (overriding any existing) a custom message banner for the organization
+      operationId: AdminService_SudoUpdateOrganizationBillingMessage
+      responses:
+        "200":
+          description: A successful response.
+          schema:
+            $ref: '#/definitions/v1SudoUpdateOrganizationBillingMessageResponse'
+        default:
+          description: An unexpected error response.
+          schema:
+            $ref: '#/definitions/rpcStatus'
+      parameters:
+        - name: org
+          in: path
+          required: true
+          type: string
+        - name: body
+          in: body
+          required: true
+          schema:
+            type: object
+            properties:
+              level:
+                $ref: '#/definitions/v1BillingIssueLevel'
+              message:
+                type: string
   /v1/superuser/projects/annotations:
     patch:
       summary: SudoUpdateAnnotations endpoint for superusers to update project annotations
@@ -4772,6 +4818,13 @@ definitions:
         $ref: '#/definitions/v1BillingIssueMetadataOnCreditTrial'
       trialCreditsDepleted:
         $ref: '#/definitions/v1BillingIssueMetadataTrialCreditsDepleted'
+      message:
+        $ref: '#/definitions/v1BillingIssueMetadataMessage'
+  v1BillingIssueMetadataMessage:
+    type: object
+    properties:
+      message:
+        type: string
   v1BillingIssueMetadataNeverSubscribed:
     type: object
   v1BillingIssueMetadataNoBillableAddress:
@@ -4867,6 +4920,7 @@ definitions:
       - BILLING_ISSUE_TYPE_NEVER_SUBSCRIBED
       - BILLING_ISSUE_TYPE_ON_CREDIT_TRIAL
       - BILLING_ISSUE_TYPE_TRIAL_CREDITS_DEPLETED
+      - BILLING_ISSUE_TYPE_MESSAGE
     default: BILLING_ISSUE_TYPE_UNSPECIFIED
   v1BillingPlan:
     type: object
@@ -6628,6 +6682,8 @@ definitions:
         format: date-time
   v1SudoDeleteOrganizationBillingIssueResponse:
     type: object
+  v1SudoDeleteOrganizationBillingMessageResponse:
+    type: object
   v1SudoGetResourceResponse:
     type: object
     properties:
@@ -6739,6 +6795,8 @@ definitions:
         $ref: '#/definitions/v1Organization'
       subscription:
         $ref: '#/definitions/v1Subscription'
+  v1SudoUpdateOrganizationBillingMessageResponse:
+    type: object
   v1SudoUpdateOrganizationCustomDomainRequest:
     type: object
     properties: