58 lines
1.4 KiB
Go
58 lines
1.4 KiB
Go
package database
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"time"
|
|
|
|
"git.toki-labs.com/toki/ariadne/internal/domain/organization"
|
|
"github.com/jackc/pgx/v5"
|
|
)
|
|
|
|
// WithinOrganization runs a transaction with PostgreSQL's tenant context set.
|
|
// Tenant-owned repositories must use this method instead of a raw pool.
|
|
func (database *Database) WithinOrganization(
|
|
ctx context.Context,
|
|
organizationID organization.ID,
|
|
operation func(pgx.Tx) error,
|
|
) error {
|
|
if database == nil || database.runtimePool == nil {
|
|
return errors.New("database is not configured")
|
|
}
|
|
if err := organizationID.Validate(); err != nil {
|
|
return fmt.Errorf("organization scope: %w", err)
|
|
}
|
|
if operation == nil {
|
|
return errors.New("organization operation is required")
|
|
}
|
|
|
|
transaction, err := database.runtimePool.BeginTx(ctx, pgx.TxOptions{})
|
|
if err != nil {
|
|
return fmt.Errorf("begin organization transaction: %w", err)
|
|
}
|
|
defer func() {
|
|
rollbackContext, cancel := context.WithTimeout(
|
|
context.Background(),
|
|
5*time.Second,
|
|
)
|
|
defer cancel()
|
|
_ = transaction.Rollback(rollbackContext)
|
|
}()
|
|
|
|
if _, err := transaction.Exec(
|
|
ctx,
|
|
"SELECT set_config('ariadne.organization_id', $1, true)",
|
|
organizationID.String(),
|
|
); err != nil {
|
|
return fmt.Errorf("set organization scope: %w", err)
|
|
}
|
|
|
|
if err := operation(transaction); err != nil {
|
|
return err
|
|
}
|
|
if err := transaction.Commit(ctx); err != nil {
|
|
return fmt.Errorf("commit organization transaction: %w", err)
|
|
}
|
|
return nil
|
|
}
|