refactor to improve database operability

jwt/generator.go

@@ -15,7 +15,7 @@ type TokenGenerator struct {
 	freshExpireAfter   int64         // Token freshness expiry time in minutes
 	trustedHost        string        // Trusted hostname to use for the tokens
 	secretKey          string        // Secret key to use for token hashing
-	db                 *sql.DB       // Database connection for token blacklisting
+	beginTx            BeginTX       // Database transaction getter for token blacklisting
 	tableConfig        TableConfig   // Table configuration
 	tableManager       *TableManager // Table lifecycle manager
 }
@@ -51,7 +51,7 @@ type GeneratorConfig struct {
 }
 
 // CreateGenerator creates and returns a new TokenGenerator using the provided configuration.
-func CreateGenerator(config GeneratorConfig) (gen *TokenGenerator, err error) {
+func CreateGenerator(config GeneratorConfig, txGetter BeginTX) (gen *TokenGenerator, err error) {
 	if config.AccessExpireAfter <= 0 {
 		return nil, errors.New("accessExpireAfter must be greater than 0")
 	}
@@ -102,7 +102,7 @@ func CreateGenerator(config GeneratorConfig) (gen *TokenGenerator, err error) {
 		freshExpireAfter:   config.FreshExpireAfter,
 		trustedHost:        config.TrustedHost,
 		secretKey:          config.SecretKey,
-		db:                 config.DB,
+		beginTx:            txGetter,
 		tableConfig:        config.TableConfig,
 		tableManager:       tableManager,
 	}, nil
@@ -112,16 +112,21 @@ func CreateGenerator(config GeneratorConfig) (gen *TokenGenerator, err error) {
 // This method should be called periodically if automatic cleanup is not enabled,
 // or can be called on-demand regardless of automatic cleanup settings.
 func (gen *TokenGenerator) Cleanup(ctx context.Context) error {
-	if gen.db == nil {
+	if gen.beginTx == nil {
 		return errors.New("No DB provided, unable to use this function")
 	}
 
+	tx, err := gen.beginTx(ctx)
+	if err != nil {
+		return pkgerrors.Wrap(err, "failed to begin transaction")
+	}
+
 	tableName := gen.tableConfig.TableName
 	currentTime := time.Now().Unix()
 
 	query := "DELETE FROM " + tableName + " WHERE exp < ?"
 
-	_, err := gen.db.ExecContext(ctx, query, currentTime)
+	_, err = tx.Exec(query, currentTime)
 	if err != nil {
 		return pkgerrors.Wrap(err, "failed to cleanup expired tokens")
 	}