235 lines
10 KiB
Go
235 lines
10 KiB
Go
package postgres
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"encoding/json"
|
|
"fmt"
|
|
|
|
"git.nianxx.cn/wangxuming/NianAIGC/backend/internal/billing"
|
|
)
|
|
|
|
const ListBillingPriceRulesSQL = `SELECT id, provider, capability, req_key, variant_key, unit, standard_unit_price_fen, markup_multiplier, enabled, conditions, quantity_source, priority, note, source, parameter_dimensions, created_at::text, updated_at::text
|
|
FROM public.billing_price_rules
|
|
WHERE ($1::boolean OR enabled = true)
|
|
ORDER BY provider, capability, id`
|
|
|
|
const GetBillingPriceRuleSQL = `SELECT id, provider, capability, req_key, variant_key, unit, standard_unit_price_fen, markup_multiplier, enabled, conditions, quantity_source, priority, note, source, parameter_dimensions, created_at::text, updated_at::text FROM public.billing_price_rules WHERE id = $1::text LIMIT 1`
|
|
const UpdateBillingPriceRuleSQL = `UPDATE public.billing_price_rules SET markup_multiplier = CASE WHEN $2::text = '' THEN $3::numeric ELSE markup_multiplier END, parameter_dimensions = CASE WHEN $2::text = '' THEN parameter_dimensions ELSE (SELECT jsonb_agg(CASE WHEN dimension->>'key' = $2::text THEN jsonb_set(dimension, '{tiers}', (SELECT jsonb_agg(CASE WHEN tier->>'value' = $4::text THEN jsonb_set(tier, '{markupMultiplier}', to_jsonb($3::numeric), true) ELSE tier END) FROM jsonb_array_elements(dimension->'tiers') tier), true) ELSE dimension END) FROM jsonb_array_elements(parameter_dimensions) dimension) END, updated_at = now() WHERE id = $1::text RETURNING id, provider, capability, req_key, variant_key, unit, standard_unit_price_fen, markup_multiplier, enabled, conditions, quantity_source, priority, note, source, parameter_dimensions, created_at::text, updated_at::text`
|
|
const GetBillingWalletSQL = `SELECT $1::text, COALESCE(balance_fen, 0), COALESCE(total_recharged_fen, 0), COALESCE(total_charged_fen, 0), COALESCE(updated_at::text, '') FROM public.billing_wallets WHERE organization_id = $1::text UNION ALL SELECT $1::text, 0, 0, 0, '' WHERE NOT EXISTS (SELECT 1 FROM public.billing_wallets WHERE organization_id = $1::text) LIMIT 1`
|
|
const ListBillingWalletsSQL = `SELECT organization_id, balance_fen, total_recharged_fen, total_charged_fen, updated_at::text FROM public.billing_wallets ORDER BY updated_at DESC`
|
|
const ListBillingLedgerSQL = `SELECT id, organization_id, COALESCE(account_id, ''), COALESCE(job_id, ''), kind, delta_fen, balance_after_fen, currency, idempotency_key, description, metadata, created_at::text FROM public.billing_ledger WHERE ($1::text = '' OR organization_id = $1::text) AND ($2::text = '' OR account_id = $2::text) ORDER BY created_at DESC LIMIT $3::integer`
|
|
const ListBillingOrganizationsSQL = `SELECT id, name, status, archive_owner_id, created_at::text, updated_at::text FROM public.platform_organizations ORDER BY created_at ASC`
|
|
const ListBillingMembersSQL = `SELECT id, display_name, phone, role, COALESCE(organization_id, ''), status FROM public.platform_users ORDER BY created_at ASC`
|
|
const BillingOrganizationExistsSQL = `SELECT EXISTS (SELECT 1 FROM public.platform_organizations WHERE id = $1::text)`
|
|
|
|
type BillingWalletPoster struct{ database *Database }
|
|
|
|
func NewBillingWalletPoster(database *Database) BillingWalletPoster {
|
|
return BillingWalletPoster{database: database}
|
|
}
|
|
func (poster BillingWalletPoster) PostWalletEntry(ctx context.Context, p billing.WalletPostParams) (billing.WalletPosting, error) {
|
|
metadata, err := json.Marshal(p.Metadata)
|
|
if err != nil {
|
|
return billing.WalletPosting{}, fmt.Errorf("encode wallet metadata: %w", err)
|
|
}
|
|
row, err := poster.database.PostWalletEntry(ctx, WalletEntryParams{LedgerID: p.LedgerID, OrganizationID: p.OrganizationID, AccountID: p.AccountID, JobID: p.JobID, Kind: p.Kind, DeltaFen: p.DeltaFen, Currency: p.Currency, IdempotencyKey: p.IdempotencyKey, Description: p.Description, Metadata: metadata})
|
|
if err != nil {
|
|
return billing.WalletPosting{}, err
|
|
}
|
|
return billing.WalletPosting{LedgerID: row.LedgerID, BalanceAfterFen: row.BalanceAfterFen, BalanceFen: row.BalanceFen, TotalRechargedFen: row.TotalRechargedFen, TotalChargedFen: row.TotalChargedFen, CreatedAt: row.CreatedAt, UpdatedAt: row.UpdatedAt, DeltaFen: row.DeltaFen}, nil
|
|
}
|
|
func (db *Database) PostBillingWalletEntry(ctx context.Context, p billing.WalletPostParams) (billing.WalletPosting, error) {
|
|
return NewBillingWalletPoster(db).PostWalletEntry(ctx, p)
|
|
}
|
|
|
|
func (db *Database) ListBillingPriceRules(ctx context.Context, includeDisabled bool) ([]billing.PriceRule, 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, ListBillingPriceRulesSQL, includeDisabled)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("list billing price rules: %w", err)
|
|
}
|
|
defer rows.Close()
|
|
var out []billing.PriceRule
|
|
for rows.Next() {
|
|
rule, err := scanFullPriceRule(rows)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
out = append(out, rule)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
|
|
func (db *Database) GetBillingPriceRule(ctx context.Context, id string) (*billing.PriceRule, error) {
|
|
rows, err := db.billingQuery(ctx, GetBillingPriceRuleSQL, id)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
if !rows.Next() {
|
|
return nil, rows.Err()
|
|
}
|
|
rule, err := scanFullPriceRule(rows)
|
|
return &rule, err
|
|
}
|
|
func (db *Database) UpdateBillingPriceRule(ctx context.Context, id string, patch billing.PricePatch) (*billing.PriceRule, error) {
|
|
rows, err := db.billingQuery(ctx, UpdateBillingPriceRuleSQL, id, patch.DimensionKey, patch.MarkupMultiplier, patch.TierValue)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
if !rows.Next() {
|
|
return nil, rows.Err()
|
|
}
|
|
rule, err := scanFullPriceRule(rows)
|
|
return &rule, err
|
|
}
|
|
func scanFullPriceRule(rows Rows) (billing.PriceRule, error) {
|
|
var rule billing.PriceRule
|
|
var req, variant, quantity, note sql.NullString
|
|
var conditions, source, dimensions json.RawMessage
|
|
if err := rows.Scan(&rule.ID, &rule.Provider, &rule.Capability, &req, &variant, &rule.Unit, &rule.StandardUnitPriceFen, &rule.MarkupMultiplier, &rule.Enabled, &conditions, &quantity, &rule.Priority, ¬e, &source, &dimensions, &rule.CreatedAt, &rule.UpdatedAt); err != nil {
|
|
return rule, fmt.Errorf("scan billing price rule: %w", err)
|
|
}
|
|
rule.ReqKey, rule.VariantKey, rule.QuantitySource, rule.Note = req.String, variant.String, billing.QuantitySource(quantity.String), note.String
|
|
if len(conditions) > 0 {
|
|
if err := json.Unmarshal(conditions, &rule.Conditions); err != nil {
|
|
return rule, err
|
|
}
|
|
}
|
|
if len(source) > 0 && string(source) != "null" {
|
|
if err := json.Unmarshal(source, &rule.Source); err != nil {
|
|
return rule, err
|
|
}
|
|
}
|
|
if len(dimensions) > 0 {
|
|
if err := json.Unmarshal(dimensions, &rule.Dimensions); err != nil {
|
|
return rule, err
|
|
}
|
|
}
|
|
return rule, nil
|
|
}
|
|
func (db *Database) BillingWallet(ctx context.Context, organizationID string) (billing.Wallet, error) {
|
|
rows, err := db.billingQuery(ctx, GetBillingWalletSQL, organizationID)
|
|
if err != nil {
|
|
return billing.Wallet{}, err
|
|
}
|
|
defer rows.Close()
|
|
if !rows.Next() {
|
|
return billing.Wallet{}, rows.Err()
|
|
}
|
|
return scanBillingWallet(rows)
|
|
}
|
|
func (db *Database) BillingWallets(ctx context.Context) ([]billing.Wallet, error) {
|
|
rows, err := db.billingQuery(ctx, ListBillingWalletsSQL)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var out []billing.Wallet
|
|
for rows.Next() {
|
|
wallet, err := scanBillingWallet(rows)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
out = append(out, wallet)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
func scanBillingWallet(rows Rows) (billing.Wallet, error) {
|
|
var wallet billing.Wallet
|
|
if err := rows.Scan(&wallet.OrganizationID, &wallet.BalanceFen, &wallet.TotalRechargedFen, &wallet.TotalChargedFen, &wallet.UpdatedAt); err != nil {
|
|
return wallet, fmt.Errorf("scan billing wallet: %w", err)
|
|
}
|
|
wallet.Currency = billing.CurrencyCNY
|
|
return wallet, nil
|
|
}
|
|
func (db *Database) BillingLedger(ctx context.Context, organizationID, accountID string, limit int) ([]billing.LedgerEntry, error) {
|
|
if limit <= 0 || limit > 500 {
|
|
limit = 500
|
|
}
|
|
rows, err := db.billingQuery(ctx, ListBillingLedgerSQL, organizationID, accountID, limit)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var out []billing.LedgerEntry
|
|
for rows.Next() {
|
|
var entry billing.LedgerEntry
|
|
var metadata json.RawMessage
|
|
if err := rows.Scan(&entry.ID, &entry.OrganizationID, &entry.AccountID, &entry.JobID, &entry.Kind, &entry.DeltaFen, &entry.BalanceAfterFen, &entry.Currency, &entry.IdempotencyKey, &entry.Description, &metadata, &entry.CreatedAt); err != nil {
|
|
return nil, err
|
|
}
|
|
if len(metadata) > 0 {
|
|
if err := json.Unmarshal(metadata, &entry.Metadata); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
out = append(out, entry)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
func (db *Database) BillingOrganizations(ctx context.Context) ([]billing.Organization, error) {
|
|
rows, err := db.billingQuery(ctx, ListBillingOrganizationsSQL)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var out []billing.Organization
|
|
for rows.Next() {
|
|
var item billing.Organization
|
|
if err := rows.Scan(&item.ID, &item.Name, &item.Status, &item.ArchiveOwnerID, &item.CreatedAt, &item.UpdatedAt); err != nil {
|
|
return nil, err
|
|
}
|
|
out = append(out, item)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
func (db *Database) BillingMembers(ctx context.Context) ([]billing.Member, error) {
|
|
rows, err := db.billingQuery(ctx, ListBillingMembersSQL)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
var out []billing.Member
|
|
for rows.Next() {
|
|
var item billing.Member
|
|
if err := rows.Scan(&item.ID, &item.DisplayName, &item.Phone, &item.Role, &item.OrganizationID, &item.Status); err != nil {
|
|
return nil, err
|
|
}
|
|
out = append(out, item)
|
|
}
|
|
return out, rows.Err()
|
|
}
|
|
func (db *Database) BillingOrganizationExists(ctx context.Context, id string) (bool, error) {
|
|
rows, err := db.billingQuery(ctx, BillingOrganizationExistsSQL, id)
|
|
if err != nil {
|
|
return false, err
|
|
}
|
|
defer rows.Close()
|
|
if !rows.Next() {
|
|
return false, rows.Err()
|
|
}
|
|
var exists bool
|
|
if err := rows.Scan(&exists); err != nil {
|
|
return false, err
|
|
}
|
|
return exists, rows.Err()
|
|
}
|
|
func (db *Database) billingQuery(ctx context.Context, query string, args ...any) (Rows, 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, query, args...)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("billing query: %w", err)
|
|
}
|
|
return rows, nil
|
|
}
|
|
|
|
var _ billing.WalletPoster = BillingWalletPoster{}
|
|
var _ billing.Store = (*Database)(nil)
|