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 }