archivmail/internal/tenantstore/quota.go

package tenantstore

import (
	"context"
	"fmt"
)

// TenantQuota holds the configurable limits for a tenant (nil = unlimited).
type TenantQuota struct {
	MaxStorageBytes *int64 `json:"max_storage_bytes"`
	MaxUsers        *int   `json:"max_users"`
	MaxEmails       *int64 `json:"max_emails"`
}

// TenantUsage holds the current resource usage for a tenant.
type TenantUsage struct {
	StorageBytes int64 `json:"storage_bytes"`
	UserCount    int64 `json:"user_count"`
	EmailCount   int64 `json:"email_count"`
}

// SetQuota updates the quota fields for a tenant.
// Pass nil to remove a limit (unlimited).
func (s *Store) SetQuota(ctx context.Context, tenantID int64, q TenantQuota) error {
	tag, err := s.pool.Exec(ctx,
		`UPDATE tenants SET max_storage_bytes=$1, max_users=$2, max_emails=$3 WHERE id=$4`,
		q.MaxStorageBytes, q.MaxUsers, q.MaxEmails, tenantID,
	)
	if err != nil {
		return fmt.Errorf("tenantstore: set quota: %w", err)
	}
	if tag.RowsAffected() == 0 {
		return fmt.Errorf("tenantstore: tenant %d not found", tenantID)
	}
	return nil
}

// GetQuota returns the quota configuration for a tenant.
func (s *Store) GetQuota(ctx context.Context, tenantID int64) (TenantQuota, error) {
	var q TenantQuota
	err := s.pool.QueryRow(ctx,
		`SELECT max_storage_bytes, max_users, max_emails FROM tenants WHERE id=$1`, tenantID,
	).Scan(&q.MaxStorageBytes, &q.MaxUsers, &q.MaxEmails)
	if err != nil {
		return q, fmt.Errorf("tenantstore: get quota: %w", err)
	}
	return q, nil
}

// GetUsage returns the current resource usage for a tenant.
func (s *Store) GetUsage(ctx context.Context, tenantID int64) (*TenantUsage, error) {
	u := &TenantUsage{}

	// Email count and storage bytes from emails table
	err := s.pool.QueryRow(ctx, `
		SELECT COUNT(*), COALESCE(SUM(size), 0)
		FROM emails
		WHERE tenant_id = $1
	`, tenantID).Scan(&u.EmailCount, &u.StorageBytes)
	if err != nil {
		return nil, fmt.Errorf("tenantstore: get usage emails: %w", err)
	}

	// User count from users table
	err = s.pool.QueryRow(ctx,
		`SELECT COUNT(*) FROM users WHERE tenant_id = $1`, tenantID,
	).Scan(&u.UserCount)
	if err != nil {
		return nil, fmt.Errorf("tenantstore: get usage users: %w", err)
	}

	return u, nil
}