93 lines
4.5 KiB
Go
93 lines
4.5 KiB
Go
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)
|
|
}
|
|
func (repository UsageRepository) ListContext(ctx context.Context, filters usage.Filters) ([]usage.Event, error) {
|
|
return repository.database.ListUsageEvents(ctx, filters)
|
|
}
|
|
|
|
var _ usage.Repository = UsageRepository{}
|