f834967768
All checks were successful
Deploy Smart Search Backend Test / deploy (push) Successful in 4m35s
109 lines
3.5 KiB
Go
109 lines
3.5 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
"strconv"
|
|
|
|
"git.techease.ru/Smart-search/smart-search-back/internal/model"
|
|
errs "git.techease.ru/Smart-search/smart-search-back/pkg/errors"
|
|
sq "github.com/Masterminds/squirrel"
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/jackc/pgx/v5/pgxpool"
|
|
)
|
|
|
|
type tokenUsageRepository struct {
|
|
pool *pgxpool.Pool
|
|
qb sq.StatementBuilderType
|
|
}
|
|
|
|
func NewTokenUsageRepository(pool *pgxpool.Pool) TokenUsageRepository {
|
|
return &tokenUsageRepository{
|
|
pool: pool,
|
|
qb: sq.StatementBuilder.PlaceholderFormat(sq.Dollar),
|
|
}
|
|
}
|
|
|
|
func (r *tokenUsageRepository) Create(ctx context.Context, usage *model.TokenUsage) error {
|
|
return r.createWithExecutor(ctx, r.pool, usage)
|
|
}
|
|
|
|
func (r *tokenUsageRepository) CreateTx(ctx context.Context, tx pgx.Tx, usage *model.TokenUsage) error {
|
|
return r.createWithExecutor(ctx, tx, usage)
|
|
}
|
|
|
|
func (r *tokenUsageRepository) createWithExecutor(ctx context.Context, exec DBTX, usage *model.TokenUsage) error {
|
|
query := r.qb.Insert("request_token_usage").Columns(
|
|
"request_id", "request_token_count", "response_token_count", "token_cost", "type",
|
|
).Values(
|
|
usage.RequestID, usage.RequestTokenCount, usage.ResponseTokenCount, usage.TokenCost, usage.Type,
|
|
).Suffix("RETURNING id, created_at")
|
|
|
|
sqlQuery, args, err := query.ToSql()
|
|
if err != nil {
|
|
return errs.NewInternalError(errs.DatabaseError, "failed to build query", err)
|
|
}
|
|
|
|
err = exec.QueryRow(ctx, sqlQuery, args...).Scan(&usage.ID, &usage.CreatedAt)
|
|
if err != nil {
|
|
return errs.NewInternalError(errs.DatabaseError, "failed to create token usage", err)
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (r *tokenUsageRepository) GetBalanceStatistics(ctx context.Context, userID int) (float64, []*model.WriteOffHistory, error) {
|
|
avgQuery := r.qb.Select("ROUND(COALESCE(AVG(COALESCE(rtu.token_cost, 0)), 0)::numeric, 2)").
|
|
From("request_token_usage rtu").
|
|
Join("requests_for_suppliers rfs ON rtu.request_id = rfs.id").
|
|
Where(sq.Eq{"rfs.user_id": userID})
|
|
|
|
avgSQL, avgArgs, err := avgQuery.ToSql()
|
|
if err != nil {
|
|
return 0, nil, errs.NewInternalError(errs.DatabaseError, "failed to build average query", err)
|
|
}
|
|
|
|
var averageCost float64
|
|
if err := r.pool.QueryRow(ctx, avgSQL, avgArgs...).Scan(&averageCost); err != nil {
|
|
return 0, nil, errs.NewInternalError(errs.DatabaseError, "failed to get average cost", err)
|
|
}
|
|
|
|
historyQuery := r.qb.Select(
|
|
"rtu.id",
|
|
"TO_CHAR(rtu.created_at, 'DD-MM-YYYY')",
|
|
"ROUND(COALESCE(rtu.token_cost, 0)::numeric, 2)",
|
|
).
|
|
From("request_token_usage rtu").
|
|
Join("requests_for_suppliers rfs ON rtu.request_id = rfs.id").
|
|
Where(sq.Eq{"rfs.user_id": userID}).
|
|
OrderBy("rtu.created_at DESC").
|
|
Limit(8)
|
|
|
|
historySQL, historyArgs, err := historyQuery.ToSql()
|
|
if err != nil {
|
|
return 0, nil, errs.NewInternalError(errs.DatabaseError, "failed to build history query", err)
|
|
}
|
|
|
|
rows, err := r.pool.Query(ctx, historySQL, historyArgs...)
|
|
if err != nil {
|
|
return 0, nil, errs.NewInternalError(errs.DatabaseError, "failed to get write-off history", err)
|
|
}
|
|
defer rows.Close()
|
|
|
|
var history []*model.WriteOffHistory
|
|
for rows.Next() {
|
|
var operationID int
|
|
var item model.WriteOffHistory
|
|
if err := rows.Scan(&operationID, &item.Data, &item.Amount); err != nil {
|
|
return 0, nil, errs.NewInternalError(errs.DatabaseError, "failed to scan write-off history", err)
|
|
}
|
|
item.OperationID = strconv.Itoa(operationID)
|
|
history = append(history, &item)
|
|
}
|
|
|
|
if err := rows.Err(); err != nil {
|
|
return 0, nil, errs.NewInternalError(errs.DatabaseError, "failed to iterate write-off history", err)
|
|
}
|
|
|
|
return averageCost, history, nil
|
|
}
|