Files
zarinpal-analytics-backend/internal/growthreport/repository.go

300 lines
14 KiB
Go

package growthreport
import (
"context"
"encoding/json"
"errors"
"fmt"
"time"
"github.com/jackc/pgx/v5"
"github.com/jackc/pgx/v5/pgxpool"
)
var (
ErrNotFound = errors.New("growth report: not found")
ErrUnknownTerminal = errors.New("growth report: unknown terminal")
ErrIdempotencyConflict = errors.New("growth report: idempotency conflict")
ErrRefreshLimited = errors.New("growth report: refresh limited")
)
type Repository struct{ pool *pgxpool.Pool }
const reportLeaseDuration = 3 * time.Minute
func NewRepository(pool *pgxpool.Pool) *Repository { return &Repository{pool: pool} }
type CreateRecord struct {
ID string
MerchantID int64
TerminalID string
Fingerprint string
DatasetWatermark string
ForecastStartDate *time.Time
CatalogVersion string
AnalyticsVersion string
PromptVersion string
SchemaVersion string
Model string
ManualRefresh bool
RefreshDay *time.Time
IdempotencyKey string
RequestHash string
}
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)
}
defer func() { _ = tx.Rollback(ctx) }()
lockKey := fmt.Sprintf("growth-revision:%d:%s", in.MerchantID, in.TerminalID)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1,0))`, lockKey); err != nil {
return Record{}, false, fmt.Errorf("locking growth report request: %w", err)
}
var terminalDBID *int64
if in.TerminalID != "" {
var id int64
if err := tx.QueryRow(ctx, `SELECT id FROM terminals WHERE merchant_id=$1 AND source_key=$2`, in.MerchantID, in.TerminalID).Scan(&id); errors.Is(err, pgx.ErrNoRows) {
return Record{}, false, ErrUnknownTerminal
} else if err != nil {
return Record{}, false, fmt.Errorf("validating report terminal: %w", err)
}
terminalDBID = &id
}
if in.IdempotencyKey != "" {
var requestHash, reportID string
err := tx.QueryRow(ctx, `SELECT request_hash,report_id FROM growth_report_idempotency WHERE merchant_id=$1 AND idempotency_key=$2`, in.MerchantID, in.IdempotencyKey).Scan(&requestHash, &reportID)
if err == nil {
if requestHash != in.RequestHash {
return Record{}, false, ErrIdempotencyConflict
}
record, err := getWith(ctx, tx, in.MerchantID, reportID)
if err != nil {
return Record{}, false, err
}
return record, true, nil
}
if !errors.Is(err, pgx.ErrNoRows) {
return Record{}, false, fmt.Errorf("reading report idempotency: %w", err)
}
}
if !in.ManualRefresh {
record, err := currentWith(ctx, tx, in.MerchantID, in.Fingerprint)
if err == nil {
return record, true, nil
}
if !errors.Is(err, ErrNotFound) {
return Record{}, false, err
}
}
if in.ManualRefresh && in.RefreshDay != nil {
refreshLockKey := fmt.Sprintf("growth-refresh:%d:%s", in.MerchantID, in.TerminalID)
if _, err := tx.Exec(ctx, `SELECT pg_advisory_xact_lock(hashtextextended($1,0))`, refreshLockKey); err != nil {
return Record{}, false, fmt.Errorf("locking growth report refresh: %w", err)
}
var hasAISummary bool
if err := tx.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM growth_reports WHERE merchant_id=$1 AND terminal_source_key IS NOT DISTINCT FROM $2::text AND status='READY' AND ai_response IS NOT NULL)`, in.MerchantID, nilIfEmpty(in.TerminalID)).Scan(&hasAISummary); err != nil {
return Record{}, false, fmt.Errorf("checking prior AI summary: %w", err)
}
var refreshUsed bool
if hasAISummary {
if err := tx.QueryRow(ctx, `SELECT EXISTS(SELECT 1 FROM growth_reports WHERE merchant_id=$1 AND terminal_source_key IS NOT DISTINCT FROM $2::text AND manual_refresh=true AND tehran_refresh_day=$3)`, in.MerchantID, nilIfEmpty(in.TerminalID), in.RefreshDay).Scan(&refreshUsed); err != nil {
return Record{}, false, fmt.Errorf("checking daily refresh: %w", err)
}
}
if refreshUsed {
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)
}
_, err = tx.Exec(ctx, `INSERT INTO growth_reports(id,merchant_id,terminal_id,terminal_source_key,fingerprint,revision,status,forecast_start_date,dataset_watermark,catalog_version,analytics_version,prompt_version,schema_version,model,manual_refresh,tehran_refresh_day)
VALUES($1,$2,$3,$4,$5,$6,'QUEUED',$7,$8,$9,$10,$11,$12,$13,$14,$15)`, in.ID, in.MerchantID, terminalDBID, nilIfEmpty(in.TerminalID), in.Fingerprint, revision, in.ForecastStartDate, in.DatasetWatermark, in.CatalogVersion, in.AnalyticsVersion, in.PromptVersion, in.SchemaVersion, in.Model, in.ManualRefresh, in.RefreshDay)
if err != nil {
return Record{}, false, fmt.Errorf("creating growth report: %w", err)
}
if in.IdempotencyKey != "" {
if _, err := tx.Exec(ctx, `INSERT INTO growth_report_idempotency(merchant_id,idempotency_key,request_hash,report_id) VALUES($1,$2,$3,$4)`, in.MerchantID, in.IdempotencyKey, in.RequestHash, in.ID); err != nil {
return Record{}, false, fmt.Errorf("creating report idempotency: %w", err)
}
}
record, err := getWith(ctx, tx, in.MerchantID, in.ID)
if err != nil {
return Record{}, false, err
}
if err := tx.Commit(ctx); err != nil {
return Record{}, false, fmt.Errorf("committing growth report: %w", err)
}
return record, false, nil
}
func (r *Repository) Get(ctx context.Context, merchantID int64, id string) (Record, error) {
return getWith(ctx, r.pool, merchantID, id)
}
func (r *Repository) Latest(ctx context.Context, merchantID int64, terminalID, analyticsVersion string) (Record, error) {
return scanRecord(r.pool.QueryRow(ctx, reportSelect+` WHERE merchant_id=$1 AND terminal_source_key IS NOT DISTINCT FROM $2::text AND analytics_version=$3 AND status IN ('READY','PARTIAL') ORDER BY created_at DESC LIMIT 1`, merchantID, nilIfEmpty(terminalID), analyticsVersion))
}
func (r *Repository) Claim(ctx context.Context) (Record, error) {
// PostgreSQL is the durable queue and coordination authority. SKIP LOCKED plus
// the state transition prevents two worker processes from generating one row.
tx, err := r.pool.Begin(ctx)
if err != nil {
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) {
return Record{}, ErrNotFound
}
if err != nil {
return Record{}, err
}
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)
}
return record, nil
}
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)
}
if command.RowsAffected() != 1 {
return ErrNotFound
}
return nil
}
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")
}
var ai any
if copy != nil {
ai = copy
}
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_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 {
QueryRow(context.Context, string, ...any) pgx.Row
}
func getWith(ctx context.Context, q queryer, merchantID int64, id string) (Record, error) {
return scanRecord(q.QueryRow(ctx, reportSelect+` WHERE merchant_id=$1 AND id=$2`, merchantID, id))
}
func currentWith(ctx context.Context, q queryer, merchantID int64, fingerprint string) (Record, error) {
return scanRecord(q.QueryRow(ctx, reportSelect+` WHERE merchant_id=$1 AND fingerprint=$2 AND manual_refresh=false AND status IN ('QUEUED','ANALYZING','GENERATING','READY','PARTIAL') ORDER BY created_at DESC LIMIT 1`, merchantID, fingerprint))
}
func scanRecord(row rowScanner) (Record, error) {
var r Record
var status string
var evidence []byte
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, &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
}
if err != nil {
return Record{}, fmt.Errorf("scanning growth report: %w", err)
}
r.Status = Status(status)
if providerRequestID != nil {
r.Usage.ProviderRequestID = *providerRequestID
}
if estimatedCost != nil {
r.Usage.EstimatedCost = *estimatedCost
}
if terminalSource != nil {
r.Scope.TerminalID = terminalSource
}
r.DeterministicEvidence = json.RawMessage(evidence)
r.EventsURL = "/api/v1/analytics/growth-reports/" + r.ReportID + "/events"
if len(ai) > 0 {
var copy AICopy
if json.Unmarshal(ai, &copy) == nil {
r.AICopy = &copy
}
}
return r, nil
}
func nilIfEmpty(value string) any {
if value == "" {
return nil
}
return value
}
func (r *Repository) NextRefreshAt(ctx context.Context, merchantID int64, terminalID string, now time.Time) (time.Time, error) {
var hasAISummary bool
var last *time.Time
if err := r.pool.QueryRow(ctx, `SELECT
EXISTS(SELECT 1 FROM growth_reports WHERE merchant_id=$1 AND terminal_source_key IS NOT DISTINCT FROM $2::text AND status='READY' AND ai_response IS NOT NULL),
max(created_at) FILTER (WHERE manual_refresh=true)
FROM growth_reports WHERE merchant_id=$1 AND terminal_source_key IS NOT DISTINCT FROM $2::text`, merchantID, nilIfEmpty(terminalID)).Scan(&hasAISummary, &last); err != nil {
return time.Time{}, fmt.Errorf("loading refresh quota: %w", err)
}
if !hasAISummary || last == nil {
return now.UTC(), nil
}
return NextTehranRefresh(*last), nil
}
func (r *Repository) DatasetMetadata(ctx context.Context) (string, time.Time, error) {
var sha string
var latest time.Time
err := r.pool.QueryRow(ctx, `SELECT ir.source_sha256,max(ps.created_at AT TIME ZONE 'Asia/Tehran')::date
FROM payment_sessions ps JOIN import_runs ir ON ir.id=ps.import_run_id WHERE ir.status='COMPLETED'
GROUP BY ir.id,ir.source_sha256 ORDER BY ir.completed_at DESC LIMIT 1`).Scan(&sha, &latest)
if err != nil {
return "", time.Time{}, fmt.Errorf("loading growth dataset metadata: %w", err)
}
return sha + ":" + latest.Format(time.DateOnly), latest.AddDate(0, 0, 1), nil
}
func (r *Repository) CleanupSuperseded(ctx context.Context, cutoff time.Time) (int64, error) {
command, err := r.pool.Exec(ctx, `DELETE FROM growth_reports old WHERE old.created_at<$1 AND EXISTS(SELECT 1 FROM growth_reports newer WHERE newer.merchant_id=old.merchant_id AND newer.terminal_source_key IS NOT DISTINCT FROM old.terminal_source_key AND newer.id<>old.id AND newer.revision>old.revision AND newer.status IN ('READY','PARTIAL'))`, cutoff)
if err != nil {
return 0, fmt.Errorf("cleaning superseded growth reports: %w", err)
}
return command.RowsAffected(), nil
}