starcoach-hackathon/zarinpal-analytics-backend
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(growth): harden AI jobs and cost controls

Browse Source
This commit is contained in:
Amirhossein Khalili
2026-08-21 01:48:47 +03:30
parent 86e9d7c5f9
commit a15c04cac1
13 changed files with 366 additions and 78 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
internal/growthreport/repository.go
View File

@@ -20,6 +20,8 @@ var (
type Repository struct{ pool *pgxpool.Pool }
const reportLeaseDuration = 3 * time.Minute
func NewRepository(pool *pgxpool.Pool) *Repository { return &Repository{pool: pool} }
type CreateRecord struct {
@@ -40,7 +42,7 @@ type CreateRecord struct {
RequestHash string
}
func (r *Repository) CreateOrReuse(ctx context.Context, in CreateRecord) (Record, bool, error) {
func (r *Repository) CreateOrReuse(ctx context.Context, in CreateRecord, beforeInsert func() error) (Record, bool, error) {
tx, err := r.pool.BeginTx(ctx, pgx.TxOptions{IsoLevel: pgx.Serializable})
if err != nil {
return Record{}, false, fmt.Errorf("beginning report transaction: %w", err)
@@ -96,6 +98,11 @@ func (r *Repository) CreateOrReuse(ctx context.Context, in CreateRecord) (Record
return Record{}, false, ErrRefreshLimited
}
}
if beforeInsert != nil {
if err := beforeInsert(); err != nil {
return Record{}, false, err
}
}
var revision int
if err := tx.QueryRow(ctx, `SELECT coalesce(max(revision),0)+1 FROM growth_reports WHERE merchant_id=$1 AND terminal_source_key IS NOT DISTINCT FROM $2::text`, in.MerchantID, nilIfEmpty(in.TerminalID)).Scan(&revision); err != nil {
return Record{}, false, fmt.Errorf("calculating growth report revision: %w", err)
@@ -139,6 +146,9 @@ func (r *Repository) Claim(ctx context.Context) (Record, error) {
return Record{}, fmt.Errorf("beginning report claim: %w", err)
}
defer func() { _ = tx.Rollback(ctx) }()
if _, err := tx.Exec(ctx, `UPDATE growth_reports SET status='QUEUED',lease_token=NULL,lease_expires_at=NULL,updated_at=now(),failure_code='WORKER_LEASE_EXPIRED' WHERE status IN ('ANALYZING','GENERATING') AND (lease_expires_at IS NULL OR lease_expires_at<now())`); err != nil {
return Record{}, fmt.Errorf("recovering expired report leases: %w", err)
}
row := tx.QueryRow(ctx, reportSelect+` WHERE id=(SELECT id FROM growth_reports WHERE status='QUEUED' ORDER BY created_at FOR UPDATE SKIP LOCKED LIMIT 1)`)
record, err := scanRecord(row)
if errors.Is(err, ErrNotFound) {
@@ -147,10 +157,13 @@ func (r *Repository) Claim(ctx context.Context) (Record, error) {
if err != nil {
return Record{}, err
}
if _, err := tx.Exec(ctx, `UPDATE growth_reports SET status='ANALYZING',updated_at=now() WHERE id=$1`, record.ReportID); err != nil {
claimToken := newReportID()
leaseExpiresAt := time.Now().UTC().Add(reportLeaseDuration)
if _, err := tx.Exec(ctx, `UPDATE growth_reports SET status='ANALYZING',lease_token=$2,lease_expires_at=$3,updated_at=now() WHERE id=$1`, record.ReportID, claimToken, leaseExpiresAt); err != nil {
return Record{}, fmt.Errorf("claiming growth report: %w", err)
}
record.Status = StatusAnalyzing
record.ClaimToken = claimToken
record.UpdatedAt = time.Now().UTC()
if err := tx.Commit(ctx); err != nil {
return Record{}, fmt.Errorf("committing report claim: %w", err)
@@ -158,8 +171,8 @@ func (r *Repository) Claim(ctx context.Context) (Record, error) {
return record, nil
}
func (r *Repository) StoreEvidence(ctx context.Context, reportID string, evidence json.RawMessage) error {
command, err := r.pool.Exec(ctx, `UPDATE growth_reports SET deterministic_evidence=$2,status='GENERATING',updated_at=now() WHERE id=$1 AND status='ANALYZING'`, reportID, evidence)
func (r *Repository) StoreEvidence(ctx context.Context, reportID, claimToken string, evidence json.RawMessage) error {
command, err := r.pool.Exec(ctx, `UPDATE growth_reports SET deterministic_evidence=$3,status='GENERATING',lease_expires_at=now()+make_interval(secs => $4),updated_at=now() WHERE id=$1 AND lease_token=$2 AND status='ANALYZING'`, reportID, claimToken, evidence, int(reportLeaseDuration/time.Second))
if err != nil {
return fmt.Errorf("storing deterministic evidence: %w", err)
}
@@ -169,7 +182,7 @@ func (r *Repository) StoreEvidence(ctx context.Context, reportID string, evidenc
return nil
}
func (r *Repository) Complete(ctx context.Context, reportID string, status Status, evidence json.RawMessage, copy *AICopy, usage Usage, failureCode string) error {
func (r *Repository) CompleteClaimed(ctx context.Context, reportID, claimToken string, status Status, evidence json.RawMessage, copy *AICopy, usage Usage, failureCode string) error {
if status != StatusReady && status != StatusPartial && status != StatusFailed {
return errors.New("growth report: invalid terminal status")
}
@@ -177,14 +190,17 @@ func (r *Repository) Complete(ctx context.Context, reportID string, status Statu
if copy != nil {
ai = copy
}
_, err := r.pool.Exec(ctx, `UPDATE growth_reports SET status=$2,deterministic_evidence=$3,ai_response=$4,failure_code=$5,input_tokens=$6,output_tokens=$7,latency_ms=$8,provider_request_id=$9,estimated_cost=$10,updated_at=now(),completed_at=now() WHERE id=$1`, reportID, status, evidence, ai, nilIfEmpty(failureCode), usage.InputTokens, usage.OutputTokens, usage.LatencyMS, nilIfEmpty(usage.ProviderRequestID), nilIfEmpty(usage.EstimatedCost))
command, err := r.pool.Exec(ctx, `UPDATE growth_reports SET status=$3,deterministic_evidence=$4,ai_response=$5,failure_code=$6,input_tokens=$7,output_tokens=$8,latency_ms=$9,provider_requests=$10,reserved_tokens=$11,provider_request_id=$12,estimated_cost=$13,lease_token=NULL,lease_expires_at=NULL,updated_at=now(),completed_at=now() WHERE id=$1 AND lease_token=$2`, reportID, claimToken, status, evidence, ai, nilIfEmpty(failureCode), usage.InputTokens, usage.OutputTokens, usage.LatencyMS, usage.ProviderRequests, usage.ReservedTokens, nilIfEmpty(usage.ProviderRequestID), nilIfEmpty(usage.EstimatedCost))
if err != nil {
return fmt.Errorf("completing growth report: %w", err)
}
if command.RowsAffected() != 1 {
return ErrNotFound
}
return nil
}
const reportSelect = `SELECT id,revision,merchant_id,terminal_id,terminal_source_key,fingerprint,status,horizon_days,forecast_start_date,dataset_watermark,catalog_version,analytics_version,prompt_version,schema_version,model,deterministic_evidence,ai_response,failure_code,input_tokens,output_tokens,latency_ms,provider_request_id,estimated_cost,manual_refresh,tehran_refresh_day,created_at,updated_at FROM growth_reports`
const reportSelect = `SELECT id,revision,merchant_id,terminal_id,terminal_source_key,fingerprint,status,horizon_days,forecast_start_date,dataset_watermark,catalog_version,analytics_version,prompt_version,schema_version,model,deterministic_evidence,ai_response,failure_code,input_tokens,output_tokens,latency_ms,provider_requests,reserved_tokens,provider_request_id,estimated_cost,manual_refresh,tehran_refresh_day,created_at,updated_at FROM growth_reports`
type rowScanner interface{ Scan(...any) error }
type queryer interface {
@@ -205,7 +221,7 @@ func scanRecord(row rowScanner) (Record, error) {
var ai []byte
var terminalSource *string
var providerRequestID, estimatedCost *string
err := row.Scan(&r.ReportID, &r.Revision, &r.MerchantID, &r.TerminalDBID, &terminalSource, &r.Fingerprint, &status, &r.HorizonDays, &r.ForecastStartDate, &r.DatasetWatermark, &r.CatalogVersion, &r.AnalyticsVersion, &r.PromptVersion, &r.SchemaVersion, &r.Model, &evidence, &ai, &r.FailureCode, &r.Usage.InputTokens, &r.Usage.OutputTokens, &r.Usage.LatencyMS, &providerRequestID, &estimatedCost, &r.ManualRefresh, &r.TehranRefreshDay, &r.CreatedAt, &r.UpdatedAt)
err := row.Scan(&r.ReportID, &r.Revision, &r.MerchantID, &r.TerminalDBID, &terminalSource, &r.Fingerprint, &status, &r.HorizonDays, &r.ForecastStartDate, &r.DatasetWatermark, &r.CatalogVersion, &r.AnalyticsVersion, &r.PromptVersion, &r.SchemaVersion, &r.Model, &evidence, &ai, &r.FailureCode, &r.Usage.InputTokens, &r.Usage.OutputTokens, &r.Usage.LatencyMS, &r.Usage.ProviderRequests, &r.Usage.ReservedTokens, &providerRequestID, &estimatedCost, &r.ManualRefresh, &r.TehranRefreshDay, &r.CreatedAt, &r.UpdatedAt)
if errors.Is(err, pgx.ErrNoRows) {
return Record{}, ErrNotFound
}