Merge pull request #234 from Flusinerd/master

feat: add tableName config for all table based stores

Author: alexedwards

cockroachdbstore/cockroachdbstore.go

@@ -2,39 +2,68 @@ package cockroachdbstore
 
 import (
 	"database/sql"
+	"fmt"
 	"log"
 	"time"
 )
 
+type Config struct {
+	// CleanUpInterval is the interval between each cleanup operation.
+	// If set to 0, the cleanup operation is disabled.
+	CleanUpInterval time.Duration
+
+	// TableName is the name of the table where the session data will be stored.
+	// If not set, it will default to "sessions".
+	TableName string
+}
+
 // CockroachDBStore represents the session store.
 type CockroachDBStore struct {
 	db          *sql.DB
 	stopCleanup chan bool
+	tableName   string
 }
 
 // New returns a new CockroachDBStore instance, with a background cleanup goroutine
 // that runs every 5 minutes to remove expired session data.
 func New(db *sql.DB) *CockroachDBStore {
-	return NewWithCleanupInterval(db, 5*time.Minute)
+	return NewWithConfig(db, Config{CleanUpInterval: 5 * time.Minute})
 }
 
 // NewWithCleanupInterval returns a new CockroachDBStore instance. The cleanupInterval
 // parameter controls how frequently expired session data is removed by the
 // background cleanup goroutine. Setting it to 0 prevents the cleanup goroutine
 // from running (i.e. expired sessions will not be removed).
 func NewWithCleanupInterval(db *sql.DB, cleanupInterval time.Duration) *CockroachDBStore {
-	p := &CockroachDBStore{db: db}
-	if cleanupInterval > 0 {
-		go p.startCleanup(cleanupInterval)
+	return NewWithConfig(db, Config{CleanUpInterval: cleanupInterval})
+}
+
+// NewWithConfig returns a new CockroachDBStore instance using the provided config.
+// If the TableName field is empty, it will be set to "sessions".
+// If the CleanUpInterval field is 0, the cleanup goroutine will not be started.
+func NewWithConfig(db *sql.DB, config Config) *CockroachDBStore {
+	if config.TableName == "" {
+		config.TableName = "sessions"
+	}
+
+	store := &CockroachDBStore{
+		db:        db,
+		tableName: config.TableName,
 	}
-	return p
+
+	if config.CleanUpInterval > 0 {
+		go store.startCleanup(config.CleanUpInterval)
+	}
+
+	return store
 }
 
 // Find returns the data for a given session token from the CockroachDBStore instance.
 // If the session token is not found or is expired, the returned exists flag will
 // be set to false.
 func (p *CockroachDBStore) Find(token string) (b []byte, exists bool, err error) {
-	row := p.db.QueryRow("SELECT data FROM sessions WHERE token = $1 AND current_timestamp < expiry", token)
+	query := fmt.Sprintf("SELECT data FROM %s WHERE token = $1 AND current_timestamp < expiry", p.tableName)
+	row := p.db.QueryRow(query, token)
 	err = row.Scan(&b)
 	if err == sql.ErrNoRows {
 		return nil, false, nil
@@ -48,24 +77,24 @@ func (p *CockroachDBStore) Find(token string) (b []byte, exists bool, err error)
 // given expiry time. If the session token already exists, then the data and expiry
 // time are updated.
 func (p *CockroachDBStore) Commit(token string, b []byte, expiry time.Time) error {
-	_, err := p.db.Exec("INSERT INTO sessions (token, data, expiry) VALUES ($1, $2, $3) ON CONFLICT (token) DO UPDATE SET data = EXCLUDED.data, expiry = EXCLUDED.expiry", token, b, expiry)
-	if err != nil {
-		return err
-	}
-	return nil
+	query := fmt.Sprintf("INSERT INTO %s (token, data, expiry) VALUES ($1, $2, $3) ON CONFLICT (token) DO UPDATE SET data = EXCLUDED.data, expiry = EXCLUDED.expiry", p.tableName)
+	_, err := p.db.Exec(query, token, b, expiry)
+	return err
 }
 
 // Delete removes a session token and corresponding data from the CockroachDBStore
 // instance.
 func (p *CockroachDBStore) Delete(token string) error {
-	_, err := p.db.Exec("DELETE FROM sessions WHERE token = $1", token)
+	query := fmt.Sprintf("DELETE FROM %s WHERE token = $1", p.tableName)
+	_, err := p.db.Exec(query, token)
 	return err
 }
 
 // All returns a map containing the token and data for all active (i.e.
 // not expired) sessions in the CockroachDBStore instance.
 func (p *CockroachDBStore) All() (map[string][]byte, error) {
-	rows, err := p.db.Query("SELECT token, data FROM sessions WHERE current_timestamp < expiry")
+	query := fmt.Sprintf("SELECT token, data FROM %s WHERE current_timestamp < expiry", p.tableName)
+	rows, err := p.db.Query(query)
 	if err != nil {
 		return nil, err
 	}
@@ -129,6 +158,7 @@ func (p *CockroachDBStore) StopCleanup() {
 }
 
 func (p *CockroachDBStore) deleteExpired() error {
-	_, err := p.db.Exec("DELETE FROM sessions WHERE expiry < current_timestamp")
+	query := fmt.Sprintf("DELETE FROM %s WHERE expiry < current_timestamp", p.tableName)
+	_, err := p.db.Exec(query)
 	return err
 }

mongodbstore/mongodbstore.go

@@ -22,25 +22,49 @@ type MongoDBStore struct {
 	stopCleanup chan bool
 }
 
+type Config struct {
+	// CleanUpInterval is the interval between each cleanup operation.
+	// If set to 0, the cleanup operation is disabled.
+	CleanupInterval time.Duration
+
+	// CollectionName is the name of the collection where the session data will be stored.
+	// If not set, it will default to "sessions".
+	CollectionName string
+}
+
 // New returns a new MongoDBStore instance, with a background cleanup goroutine that
 // runs every minute to remove expired session data.
 func New(db *mongo.Database) *MongoDBStore {
-	return NewWithCleanupInterval(db, time.Minute)
+	return NewWithConfig(db, Config{
+		CleanupInterval: time.Minute,
+	})
 }
 
 // NewWithCleanupInterval returns a new MongoDBStore instance. The cleanupInterval
 // parameter controls how frequently expired session data is removed by the
 // background cleanup goroutine. Setting it to 0 prevents the cleanup goroutine
 // from running (i.e. expired sessions will not be removed).
 func NewWithCleanupInterval(db *mongo.Database, cleanupInterval time.Duration) *MongoDBStore {
-	collection := db.Collection("sessions")
+	return NewWithConfig(db, Config{
+		CleanupInterval: cleanupInterval,
+	})
+}
+
+// NewWithConfig returns a new MongoDBStore instance with the given configuration.
+// If the CollectionName field is empty, it will be set to "sessions".
+// If the CleanupInterval field is 0, the cleanup goroutine will not be started.
+func NewWithConfig(db *mongo.Database, config Config) *MongoDBStore {
+	if config.CollectionName == "" {
+		config.CollectionName = "sessions"
+	}
+	collection := db.Collection(config.CollectionName)
 
 	m := &MongoDBStore{
 		collection: collection,
 	}
 
-	if cleanupInterval > 0 {
-		go m.startCleanup(cleanupInterval)
+	if config.CleanupInterval > 0 {
+		go m.startCleanup(config.CleanupInterval)
 	}
 
 	return m

mssqlstore/mssqlstore.go

@@ -2,6 +2,7 @@ package mssqlstore
 
 import (
 	"database/sql"
+	"fmt"
 	"log"
 	"time"
 )
@@ -10,31 +11,59 @@ import (
 type MSSQLStore struct {
 	db          *sql.DB
 	stopCleanup chan bool
+	tableName   string
+}
+
+type Config struct {
+	// CleanUpInterval is the interval between each cleanup operation.
+	// If set to 0, the cleanup operation is disabled.
+	CleanUpInterval time.Duration
+
+	// TableName is the name of the table where the session data will be stored.
+	// If not set, it will default to "sessions".
+	TableName string
 }
 
 // New returns a new MSSQLStore instance, with a background cleanup goroutine
 // that runs every 5 minutes to remove expired session data.
 func New(db *sql.DB) *MSSQLStore {
-	return NewWithCleanupInterval(db, 5*time.Minute)
+	return NewWithConfig(db, Config{
+		CleanUpInterval: 5 * time.Minute,
+	})
 }
 
 // NewWithCleanupInterval returns a new MSSQLStore instance. The cleanupInterval
 // parameter controls how frequently expired session data is removed by the
 // background cleanup goroutine. Setting it to 0 prevents the cleanup goroutine
 // from running (i.e. expired sessions will not be removed).
 func NewWithCleanupInterval(db *sql.DB, cleanupInterval time.Duration) *MSSQLStore {
-	m := &MSSQLStore{db: db}
-	if cleanupInterval > 0 {
-		go m.startCleanup(cleanupInterval)
+	return NewWithConfig(db, Config{
+		CleanUpInterval: cleanupInterval,
+	})
+}
+
+// NewWithConfig returns a new MSSQLStore instance with the given configuration.
+// If the TableName field is empty, it will be set to "sessions".
+// If the CleanUpInterval field is 0, the cleanup goroutine will not be started.
+func NewWithConfig(db *sql.DB, config Config) *MSSQLStore {
+	if config.TableName == "" {
+		config.TableName = "sessions"
 	}
+
+	m := &MSSQLStore{db: db, tableName: config.TableName}
+	if config.CleanUpInterval > 0 {
+		go m.startCleanup(config.CleanUpInterval)
+	}
+
 	return m
 }
 
 // Find returns the data for a given session token from the MSSQLStore instance.
 // If the session token is not found or is expired, the returned exists flag will
 // be set to false.
 func (m *MSSQLStore) Find(token string) (b []byte, exists bool, err error) {
-	row := m.db.QueryRow("SELECT data FROM sessions WHERE token = @p1 AND GETUTCDATE() < expiry", token)
+	query := fmt.Sprintf("SELECT data FROM %s WHERE token = @p1 AND GETUTCDATE() < expiry", m.tableName)
+	row := m.db.QueryRow(query, token)
 	err = row.Scan(&b)
 	if err == sql.ErrNoRows {
 		return nil, false, nil
@@ -48,26 +77,26 @@ func (m *MSSQLStore) Find(token string) (b []byte, exists bool, err error) {
 // given expiry time. If the session token already exists, then the data and expiry
 // time are updated.
 func (m *MSSQLStore) Commit(token string, b []byte, expiry time.Time) error {
-	_, err := m.db.Exec(`MERGE INTO sessions WITH (HOLDLOCK) AS T USING (VALUES(@p1)) AS S (token) ON (T.token = S.token)
-						 WHEN MATCHED THEN UPDATE SET data = @p2, expiry = @p3
-						 WHEN NOT MATCHED THEN INSERT (token, data, expiry) VALUES(@p1, @p2, @p3);`, token, b, expiry.UTC())
-	if err != nil {
-		return err
-	}
-	return nil
+	query := fmt.Sprintf(`MERGE INTO %s WITH (HOLDLOCK) AS T USING (VALUES(@p1)) AS S (token) ON (T.token = S.token)
+		WHEN MATCHED THEN UPDATE SET data = @p2, expiry = @p3
+		WHEN NOT MATCHED THEN INSERT (token, data, expiry) VALUES(@p1, @p2, @p3);`, m.tableName)
+	_, err := m.db.Exec(query, token, b, expiry.UTC())
+	return err
 }
 
 // Delete removes a session token and corresponding data from the MSSQLStore
 // instance.
 func (m *MSSQLStore) Delete(token string) error {
-	_, err := m.db.Exec("DELETE FROM sessions WHERE token = @p1", token)
+	query := fmt.Sprintf("DELETE FROM %s WHERE token = @p1", m.tableName)
+	_, err := m.db.Exec(query, token)
 	return err
 }
 
 // All returns a map containing the token and data for all active (i.e.
 // not expired) sessions in the MSSQLStore instance.
 func (m *MSSQLStore) All() (map[string][]byte, error) {
-	rows, err := m.db.Query("SELECT token, data FROM sessions WHERE GETUTCDATE() < expiry")
+	query := fmt.Sprintf("SELECT token, data FROM %s WHERE GETUTCDATE() < expiry", m.tableName)
+	rows, err := m.db.Query(query)
 	if err != nil {
 		return nil, err
 	}
@@ -115,7 +144,7 @@ func (m *MSSQLStore) startCleanup(interval time.Duration) {
 }
 
 // StopCleanup terminates the background cleanup goroutine for the MSSQLStore
-// instance. It's rare to terminate this; generally MSSQLStore instances and
+// instance. It's rare to terminate this; generally MSSQLStore instances and	// instance.
 // their cleanup goroutines are intended to be long-lived and run for the lifetime
 // of your application.
 //
@@ -131,6 +160,7 @@ func (m *MSSQLStore) StopCleanup() {
 }
 
 func (m *MSSQLStore) deleteExpired() error {
-	_, err := m.db.Exec("DELETE FROM sessions WHERE expiry < GETUTCDATE()")
+	query := fmt.Sprintf("DELETE FROM %s WHERE expiry < GETUTCDATE()", m.tableName)
+	_, err := m.db.Exec(query)
 	return err
 }