agentic-coding-workflow/internal/mcp/repository.go

package mcp

import (
	"context"
	"encoding/json"
	"errors"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
	mcpsqlc "github.com/yan1h/agent-coding-workflow/internal/mcp/sqlc"
)

// Repository 是 mcp 模块对 PG 的全部依赖。Service 单测通过手写 fakeRepo 满足。
type Repository interface {
	// CRUD
	Create(ctx context.Context, t *Token) (*Token, error)
	GetByHash(ctx context.Context, hash []byte) (*Token, error)
	GetByID(ctx context.Context, id uuid.UUID) (*Token, error)
	List(ctx context.Context, callerUserID *uuid.UUID, activeOnly bool) ([]*Token, error)

	// 状态变更
	UpdateLastUsed(ctx context.Context, id uuid.UUID) error
	Revoke(ctx context.Context, id uuid.UUID) (uuid.UUID, error)
	RevokeBySession(ctx context.Context, sessionID uuid.UUID) ([]uuid.UUID, error)
	ReapStaleSystemTokens(ctx context.Context) ([]uuid.UUID, error)
	PurgeExpired(ctx context.Context, before time.Time) (int64, error)

	// 测试辅助：直接物理删（生产不暴露）
	DeleteByID(ctx context.Context, id uuid.UUID) error

	// WithTx 返回绑定到指定事务的 Repository。详见 pgRepo.WithTx 文档。
	WithTx(tx pgx.Tx) Repository
}

// NewPostgresRepository 用 pgxpool 构造 Repository。
func NewPostgresRepository(pool *pgxpool.Pool) Repository {
	return &pgRepo{pool: pool, q: mcpsqlc.New(pool)}
}

type pgRepo struct {
	pool *pgxpool.Pool
	q    *mcpsqlc.Queries
}

// WithTx 返回一个绑定到指定事务的 Repository 副本。Caller 在外部用 pgx.BeginFunc
// 或 acp.Repository.InTx 取得 tx；这里只负责生成一个把所有 sqlc 调用切到 tx 的轻
// 实例。tx-bound 实例不持有 pool（防止误用同一实例混 pool/tx 两路）。
func (r *pgRepo) WithTx(tx pgx.Tx) Repository {
	return &pgRepo{pool: nil, q: r.q.WithTx(tx)}
}

// ===== 类型转换 helper =====

func toPgUUID(u uuid.UUID) pgtype.UUID {
	return pgtype.UUID{Bytes: u, Valid: true}
}
func toPgUUIDPtr(u *uuid.UUID) pgtype.UUID {
	if u == nil {
		return pgtype.UUID{}
	}
	return pgtype.UUID{Bytes: *u, Valid: true}
}
func fromPgUUID(p pgtype.UUID) uuid.UUID {
	if !p.Valid {
		return uuid.Nil
	}
	return uuid.UUID(p.Bytes)
}
func fromPgUUIDPtr(p pgtype.UUID) *uuid.UUID {
	if !p.Valid {
		return nil
	}
	id := uuid.UUID(p.Bytes)
	return &id
}
func toPgTimestamp(t *time.Time) pgtype.Timestamptz {
	if t == nil {
		return pgtype.Timestamptz{}
	}
	return pgtype.Timestamptz{Time: *t, Valid: true}
}
func fromPgTimestamp(t pgtype.Timestamptz) *time.Time {
	if !t.Valid {
		return nil
	}
	out := t.Time
	return &out
}

// scopeToJSON 把 Scope 序列化为 jsonb 入参。空 Scope 也写 `{}`，与 DB DEFAULT 一致。
func scopeToJSON(s Scope) ([]byte, error) {
	return json.Marshal(s)
}

// scopeFromJSON 反序列化 DB 中 jsonb 字段。空对象 `{}` → 零值 Scope（Tools/ProjectIDs 都为 nil）。
func scopeFromJSON(raw []byte) (Scope, error) {
	if len(raw) == 0 {
		return Scope{}, nil
	}
	var s Scope
	if err := json.Unmarshal(raw, &s); err != nil {
		return Scope{}, err
	}
	return s, nil
}

// rowToToken 把 sqlc 生成的 row 翻译为业务结构。
func rowToToken(r mcpsqlc.McpToken) (*Token, error) {
	scope, err := scopeFromJSON(r.Scope)
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "decode mcp_token.scope")
	}
	return &Token{
		ID:           fromPgUUID(r.ID),
		TokenHash:    r.TokenHash,
		UserID:       fromPgUUID(r.UserID),
		Name:         r.Name,
		Issuer:       Issuer(r.Issuer),
		Scope:        scope,
		ACPSessionID: fromPgUUIDPtr(r.AcpSessionID),
		ExpiresAt:    fromPgTimestamp(r.ExpiresAt),
		LastUsedAt:   fromPgTimestamp(r.LastUsedAt),
		RevokedAt:    fromPgTimestamp(r.RevokedAt),
		CreatedBy:    fromPgUUID(r.CreatedBy),
		CreatedAt:    r.CreatedAt.Time,
	}, nil
}

// pgxNoRowsToErrNotFound 把 pgx.ErrNoRows 翻译为 errs.AppError(NotFound)。
func pgxNoRowsToErrNotFound(err error, code errs.Code, msg string) error {
	if errors.Is(err, pgx.ErrNoRows) {
		return errs.New(code, msg)
	}
	return errs.Wrap(err, errs.CodeInternal, msg)
}

// ===== Create =====

func (r *pgRepo) Create(ctx context.Context, t *Token) (*Token, error) {
	scopeJSON, err := scopeToJSON(t.Scope)
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "encode mcp_token.scope")
	}
	row, err := r.q.CreateMcpToken(ctx, mcpsqlc.CreateMcpTokenParams{
		ID:           toPgUUID(t.ID),
		TokenHash:    t.TokenHash,
		UserID:       toPgUUID(t.UserID),
		Name:         t.Name,
		Issuer:       string(t.Issuer),
		Scope:        scopeJSON,
		AcpSessionID: toPgUUIDPtr(t.ACPSessionID),
		ExpiresAt:    toPgTimestamp(t.ExpiresAt),
		CreatedBy:    toPgUUID(t.CreatedBy),
	})
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "insert mcp_token")
	}
	return rowToToken(row)
}

// ===== Get =====

func (r *pgRepo) GetByHash(ctx context.Context, hash []byte) (*Token, error) {
	row, err := r.q.GetMcpTokenByHash(ctx, hash)
	if err != nil {
		return nil, pgxNoRowsToErrNotFound(err, errs.CodeMcpTokenInvalid, "mcp token not found")
	}
	return rowToToken(row)
}

func (r *pgRepo) GetByID(ctx context.Context, id uuid.UUID) (*Token, error) {
	row, err := r.q.GetMcpTokenByID(ctx, toPgUUID(id))
	if err != nil {
		return nil, pgxNoRowsToErrNotFound(err, errs.CodeMcpTokenNotFound, "mcp token not found")
	}
	return rowToToken(row)
}

// ===== Stub methods (implemented in subsequent tasks) =====

func (r *pgRepo) List(ctx context.Context, callerUserID *uuid.UUID, activeOnly bool) ([]*Token, error) {
	rows, err := r.q.ListMcpTokens(ctx, mcpsqlc.ListMcpTokensParams{
		CallerUserID: toPgUUIDPtr(callerUserID),
		ActiveOnly:   activeOnly,
	})
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "list mcp_tokens")
	}
	out := make([]*Token, 0, len(rows))
	for _, row := range rows {
		t, err := rowToToken(row)
		if err != nil {
			return nil, err
		}
		out = append(out, t)
	}
	return out, nil
}

func (r *pgRepo) UpdateLastUsed(ctx context.Context, id uuid.UUID) error {
	if err := r.q.UpdateMcpTokenLastUsed(ctx, toPgUUID(id)); err != nil {
		return errs.Wrap(err, errs.CodeInternal, "update mcp_token last_used_at")
	}
	return nil
}

func (r *pgRepo) Revoke(ctx context.Context, id uuid.UUID) (uuid.UUID, error) {
	row, err := r.q.RevokeMcpToken(ctx, toPgUUID(id))
	if err != nil {
		// pgx.ErrNoRows = 已撤销 / 不存在；统一返 NotFound 让上层细分
		return uuid.Nil, pgxNoRowsToErrNotFound(err, errs.CodeMcpTokenNotFound,
			"mcp token not found or already revoked")
	}
	return fromPgUUID(row), nil
}

func (r *pgRepo) RevokeBySession(ctx context.Context, sessionID uuid.UUID) ([]uuid.UUID, error) {
	rows, err := r.q.RevokeMcpTokensBySession(ctx, toPgUUID(sessionID))
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "revoke mcp_tokens by session")
	}
	out := make([]uuid.UUID, 0, len(rows))
	for _, row := range rows {
		out = append(out, fromPgUUID(row))
	}
	return out, nil
}

func (r *pgRepo) ReapStaleSystemTokens(ctx context.Context) ([]uuid.UUID, error) {
	rows, err := r.q.ReapStaleSystemTokens(ctx)
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "reap stale system tokens")
	}
	out := make([]uuid.UUID, 0, len(rows))
	for _, row := range rows {
		out = append(out, fromPgUUID(row))
	}
	return out, nil
}

func (r *pgRepo) PurgeExpired(ctx context.Context, before time.Time) (int64, error) {
	n, err := r.q.PurgeMcpTokensExpired(ctx, pgtype.Timestamptz{Time: before, Valid: true})
	if err != nil {
		return 0, errs.Wrap(err, errs.CodeInternal, "purge expired mcp_tokens")
	}
	return n, nil
}

func (r *pgRepo) DeleteByID(ctx context.Context, id uuid.UUID) error {
	if err := r.q.DeleteMcpTokenByID(ctx, toPgUUID(id)); err != nil {
		return errs.Wrap(err, errs.CodeInternal, "delete mcp_token")
	}
	return nil
}