package db

import (
	"context"

	"github.com/jackc/pgx/v5"
)

type ctxKey string

const txCtxKey ctxKey = "poolproxy_tx"

// ContextWithTx stores a pgx.Tx in the context for PoolProxy routing.
// When PoolProxy.Exec/Query/QueryRow sees this context key, it routes the
// call through the stored transaction instead of the pool.
func ContextWithTx(ctx context.Context, tx pgx.Tx) context.Context {
	return context.WithValue(ctx, txCtxKey, tx)
}

// TxFromContext extracts a pgx.Tx from context (returns nil if none active).
func TxFromContext(ctx context.Context) pgx.Tx {
	tx, _ := ctx.Value(txCtxKey).(pgx.Tx)
	return tx
}