52 lines
1.5 KiB
Go
52 lines
1.5 KiB
Go
package repository
|
|
|
|
import (
|
|
"context"
|
|
|
|
"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
|
|
}
|