feat: add billing and usage domain cores

backend/internal/postgres/usage.go

@@ -0,0 +1,89 @@
+package postgres
+
+import (
+	"context"
+	"database/sql"
+	"fmt"
+
+	"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/usage"
+)
+
+const InsertUsageEventSQL = `INSERT INTO public.usage_events (id, owner_id, job_id, source, capability, provider, req_key, account_username, account_display_name, tenant_id, organization_id, organization_name, quantity, estimated_unit, charged_amount_fen, currency, created_at)
+VALUES ($1::text, $2::text, $3::text, $4::text, $5::text, NULLIF($6::text, ''), NULLIF($7::text, ''), NULLIF($8::text, ''), NULLIF($9::text, ''), NULLIF($10::text, ''), NULLIF($11::text, ''), NULLIF($12::text, ''), $13::integer, $14::text, $15::bigint, NULLIF($16::text, ''), $17::timestamptz)
+ON CONFLICT (job_id) DO NOTHING
+RETURNING id, owner_id, job_id, source, capability, COALESCE(provider, ''), COALESCE(req_key, ''), COALESCE(account_username, ''), COALESCE(account_display_name, ''), COALESCE(tenant_id, ''), COALESCE(organization_id, ''), COALESCE(organization_name, ''), quantity, estimated_unit, charged_amount_fen, COALESCE(currency, ''), created_at::text`
+const ListUsageEventsSQL = `SELECT id, owner_id, job_id, source, capability, COALESCE(provider, ''), COALESCE(req_key, ''), COALESCE(account_username, ''), COALESCE(account_display_name, ''), COALESCE(tenant_id, ''), COALESCE(organization_id, ''), COALESCE(organization_name, ''), quantity, estimated_unit, charged_amount_fen, COALESCE(currency, ''), created_at::text
+FROM public.usage_events
+WHERE ($1::text = '' OR owner_id = $1::text) AND ($2::text = '' OR organization_id = $2::text) AND ($3::timestamptz IS NULL OR created_at >= $3::timestamptz) AND ($4::timestamptz IS NULL OR created_at < $4::timestamptz)
+ORDER BY created_at DESC, id DESC`
+
+func (db *Database) InsertUsageEvent(ctx context.Context, e usage.Event) (usage.Event, bool, error) {
+	if db.config.Backend != BackendPostgres || db.querier == nil {
+		return usage.Event{}, false, fmt.Errorf("PostgreSQL is unavailable when ZHINIAN_DATA_BACKEND=%s", db.config.Backend)
+	}
+	rows, err := db.querier.Query(ctx, InsertUsageEventSQL, e.ID, e.OwnerID, e.JobID, e.Source, e.Capability, e.Provider, e.ReqKey, e.AccountUsername, e.AccountDisplayName, e.TenantID, e.OrganizationID, e.OrganizationName, e.Quantity, e.EstimatedUnit, e.ChargedAmountFen, e.Currency, e.CreatedAt)
+	if err != nil {
+		return usage.Event{}, false, fmt.Errorf("insert usage event: %w", err)
+	}
+	defer rows.Close()
+	if !rows.Next() {
+		return usage.Event{}, false, rows.Err()
+	}
+	out, err := scanUsage(rows)
+	return out, err == nil, err
+}
+func (db *Database) ListUsageEvents(ctx context.Context, f usage.Filters) ([]usage.Event, error) {
+	if db.config.Backend != BackendPostgres || db.querier == nil {
+		return nil, fmt.Errorf("PostgreSQL is unavailable when ZHINIAN_DATA_BACKEND=%s", db.config.Backend)
+	}
+	rows, err := db.querier.Query(ctx, ListUsageEventsSQL, f.OwnerID, f.OrganizationID, nullableTime(f.From), nullableTime(f.To))
+	if err != nil {
+		return nil, fmt.Errorf("list usage events: %w", err)
+	}
+	defer rows.Close()
+	var out []usage.Event
+	for rows.Next() {
+		event, err := scanUsage(rows)
+		if err != nil {
+			return nil, err
+		}
+		if f.Capability != "" && event.Capability != f.Capability || f.Provider != "" && event.Provider != f.Provider {
+			continue
+		}
+		out = append(out, event)
+	}
+	return out, rows.Err()
+}
+func scanUsage(rows Rows) (usage.Event, error) {
+	var e usage.Event
+	var charged sql.NullInt64
+	if err := rows.Scan(&e.ID, &e.OwnerID, &e.JobID, &e.Source, &e.Capability, &e.Provider, &e.ReqKey, &e.AccountUsername, &e.AccountDisplayName, &e.TenantID, &e.OrganizationID, &e.OrganizationName, &e.Quantity, &e.EstimatedUnit, &charged, &e.Currency, &e.CreatedAt); err != nil {
+		return usage.Event{}, fmt.Errorf("scan usage event: %w", err)
+	}
+	if charged.Valid {
+		e.ChargedAmountFen = &charged.Int64
+	}
+	return e, nil
+}
+func nullableTime(v string) any {
+	if v == "" {
+		return nil
+	}
+	return v
+}
+
+type UsageRepository struct{ database *Database }
+
+func NewUsageRepository(database *Database) UsageRepository {
+	return UsageRepository{database: database}
+}
+
+func (repository UsageRepository) Insert(event usage.Event) (usage.Event, bool, error) {
+	return repository.database.InsertUsageEvent(context.Background(), event)
+}
+
+func (repository UsageRepository) List(filters usage.Filters) ([]usage.Event, error) {
+	return repository.database.ListUsageEvents(context.Background(), filters)
+}
+
+var _ usage.Repository = UsageRepository{}