// Copyright (C) 2024 Mya Pitzeruse
// SPDX-License-Identifier: MIT

package crud

import (
	"context"

	"gorm.io/gorm"
)

type contextKey string

const transaction contextKey = "go.pitz.tech/gorm/crud/txn"

// Transaction returns the transaction associated with the context. If none is present, the db is used.
func Transaction(ctx context.Context, db *gorm.DB) *gorm.DB {
	txn, ok := ctx.Value(transaction).(*gorm.DB)
	if ok {
		return txn
	}

	return db
}

// WithTransaction attaches the current transaction to provided context.
func WithTransaction(ctx context.Context, db *gorm.DB) context.Context {
	return context.WithValue(ctx, transaction, db)
}