Refine again

This commit is contained in:
binwiederhier
2026-03-01 14:24:06 -05:00
parent 10a6939d8e
commit 11c79a6369
6 changed files with 133 additions and 173 deletions

View File

@@ -103,7 +103,7 @@ var pgQueries = queries{
// NewPostgresStore creates a new PostgreSQL-backed message cache store using an existing database connection pool. // NewPostgresStore creates a new PostgreSQL-backed message cache store using an existing database connection pool.
func NewPostgresStore(db *sql.DB, batchSize int, batchTimeout time.Duration) (*Cache, error) { func NewPostgresStore(db *sql.DB, batchSize int, batchTimeout time.Duration) (*Cache, error) {
if err := setupPostgresDB(db); err != nil { if err := setupPostgres(db); err != nil {
return nil, err return nil, err
} }
return newCache(db, pgQueries, nil, batchSize, batchTimeout, false), nil return newCache(db, pgQueries, nil, batchSize, batchTimeout, false), nil

View File

@@ -60,7 +60,7 @@ const (
postgresSelectSchemaVersionQuery = `SELECT version FROM schema_version WHERE store = 'message'` postgresSelectSchemaVersionQuery = `SELECT version FROM schema_version WHERE store = 'message'`
) )
func setupPostgresDB(db *sql.DB) error { func setupPostgres(db *sql.DB) error {
var schemaVersion int var schemaVersion int
err := db.QueryRow(postgresSelectSchemaVersionQuery).Scan(&schemaVersion) err := db.QueryRow(postgresSelectSchemaVersionQuery).Scan(&schemaVersion)
if err != nil { if err != nil {

View File

@@ -57,19 +57,25 @@ type Manager struct {
var _ Auther = (*Manager)(nil) var _ Auther = (*Manager)(nil)
// initManager sets defaults and runs startup tasks common to all backends func newManager(db *sql.DB, queries queries, config *Config) (*Manager, error) {
func initManager(manager *Manager) error { if config.BcryptCost <= 0 {
if manager.config.BcryptCost <= 0 { config.BcryptCost = DefaultUserPasswordBcryptCost
manager.config.BcryptCost = DefaultUserPasswordBcryptCost
} }
if manager.config.QueueWriterInterval.Seconds() <= 0 { if config.QueueWriterInterval.Seconds() <= 0 {
manager.config.QueueWriterInterval = DefaultUserStatsQueueWriterInterval config.QueueWriterInterval = DefaultUserStatsQueueWriterInterval
}
manager := &Manager{
config: config,
db: db,
statsQueue: make(map[string]*Stats),
tokenQueue: make(map[string]*TokenUpdate),
queries: queries,
} }
if err := manager.maybeProvisionUsersAccessAndTokens(); err != nil { if err := manager.maybeProvisionUsersAccessAndTokens(); err != nil {
return err return nil, err
} }
go manager.asyncQueueWriter(manager.config.QueueWriterInterval) go manager.asyncQueueWriter(manager.config.QueueWriterInterval)
return nil return manager, nil
} }
// Authenticate checks username and password and returns a User if correct, and the user has not been // Authenticate checks username and password and returns a User if correct, and the user has not been

View File

@@ -204,17 +204,7 @@ const (
) )
// NewPostgresManager creates a new Manager backed by a PostgreSQL database using an existing connection pool. // NewPostgresManager creates a new Manager backed by a PostgreSQL database using an existing connection pool.
func NewPostgresManager(db *sql.DB, config *Config) (*Manager, error) { var postgresQueries = queries{
if err := setupPostgres(db); err != nil {
return nil, err
}
manager := &Manager{
config: config,
db: db,
statsQueue: make(map[string]*Stats),
tokenQueue: make(map[string]*TokenUpdate),
queries: queries{
// User queries
selectUserByID: postgresSelectUserByIDQuery, selectUserByID: postgresSelectUserByIDQuery,
selectUserByName: postgresSelectUserByNameQuery, selectUserByName: postgresSelectUserByNameQuery,
selectUserByToken: postgresSelectUserByTokenQuery, selectUserByToken: postgresSelectUserByTokenQuery,
@@ -234,8 +224,6 @@ func NewPostgresManager(db *sql.DB, config *Config) (*Manager, error) {
deleteUser: postgresDeleteUserQuery, deleteUser: postgresDeleteUserQuery,
deleteUserTier: postgresDeleteUserTierQuery, deleteUserTier: postgresDeleteUserTierQuery,
deleteUsersMarked: postgresDeleteUsersMarkedQuery, deleteUsersMarked: postgresDeleteUsersMarkedQuery,
// Access queries
selectTopicPerms: postgresSelectTopicPermsQuery, selectTopicPerms: postgresSelectTopicPermsQuery,
selectUserAllAccess: postgresSelectUserAllAccessQuery, selectUserAllAccess: postgresSelectUserAllAccessQuery,
selectUserAccess: postgresSelectUserAccessQuery, selectUserAccess: postgresSelectUserAccessQuery,
@@ -249,8 +237,6 @@ func NewPostgresManager(db *sql.DB, config *Config) (*Manager, error) {
deleteUserAccessProvisioned: postgresDeleteUserAccessProvisionedQuery, deleteUserAccessProvisioned: postgresDeleteUserAccessProvisionedQuery,
deleteTopicAccess: postgresDeleteTopicAccessQuery, deleteTopicAccess: postgresDeleteTopicAccessQuery,
deleteAllAccess: postgresDeleteAllAccessQuery, deleteAllAccess: postgresDeleteAllAccessQuery,
// Token queries
selectToken: postgresSelectTokenQuery, selectToken: postgresSelectTokenQuery,
selectTokens: postgresSelectTokensQuery, selectTokens: postgresSelectTokensQuery,
selectTokenCount: postgresSelectTokenCountQuery, selectTokenCount: postgresSelectTokenCountQuery,
@@ -263,26 +249,22 @@ func NewPostgresManager(db *sql.DB, config *Config) (*Manager, error) {
deleteAllToken: postgresDeleteAllTokenQuery, deleteAllToken: postgresDeleteAllTokenQuery,
deleteExpiredTokens: postgresDeleteExpiredTokensQuery, deleteExpiredTokens: postgresDeleteExpiredTokensQuery,
deleteExcessTokens: postgresDeleteExcessTokensQuery, deleteExcessTokens: postgresDeleteExcessTokensQuery,
// Tier queries
insertTier: postgresInsertTierQuery, insertTier: postgresInsertTierQuery,
selectTiers: postgresSelectTiersQuery, selectTiers: postgresSelectTiersQuery,
selectTierByCode: postgresSelectTierByCodeQuery, selectTierByCode: postgresSelectTierByCodeQuery,
selectTierByPriceID: postgresSelectTierByPriceIDQuery, selectTierByPriceID: postgresSelectTierByPriceIDQuery,
updateTier: postgresUpdateTierQuery, updateTier: postgresUpdateTierQuery,
deleteTier: postgresDeleteTierQuery, deleteTier: postgresDeleteTierQuery,
// Phone queries
selectPhoneNumbers: postgresSelectPhoneNumbersQuery, selectPhoneNumbers: postgresSelectPhoneNumbersQuery,
insertPhoneNumber: postgresInsertPhoneNumberQuery, insertPhoneNumber: postgresInsertPhoneNumberQuery,
deletePhoneNumber: postgresDeletePhoneNumberQuery, deletePhoneNumber: postgresDeletePhoneNumberQuery,
// Billing queries
updateBilling: postgresUpdateBillingQuery, updateBilling: postgresUpdateBillingQuery,
},
} }
if err := initManager(manager); err != nil {
// NewPostgresManager creates a new Manager backed by a PostgreSQL database
func NewPostgresManager(db *sql.DB, config *Config) (*Manager, error) {
if err := setupPostgres(db); err != nil {
return nil, err return nil, err
} }
return manager, nil return newManager(db, postgresQueries, config)
} }

View File

@@ -201,28 +201,7 @@ const (
` `
) )
// NewSQLiteManager creates a new Manager backed by a SQLite database var sqliteQueries = queries{
func NewSQLiteManager(filename, startupQueries string, config *Config) (*Manager, error) {
parentDir := filepath.Dir(filename)
if !util.FileExists(parentDir) {
return nil, fmt.Errorf("user database directory %s does not exist or is not accessible", parentDir)
}
db, err := sql.Open("sqlite3", filename)
if err != nil {
return nil, err
}
if err := setupSQLite(db); err != nil {
return nil, err
}
if err := runSQLiteStartupQueries(db, startupQueries); err != nil {
return nil, err
}
manager := &Manager{
config: config,
db: db,
statsQueue: make(map[string]*Stats),
tokenQueue: make(map[string]*TokenUpdate),
queries: queries{
selectUserByID: sqliteSelectUserByIDQuery, selectUserByID: sqliteSelectUserByIDQuery,
selectUserByName: sqliteSelectUserByNameQuery, selectUserByName: sqliteSelectUserByNameQuery,
selectUserByToken: sqliteSelectUserByTokenQuery, selectUserByToken: sqliteSelectUserByTokenQuery,
@@ -277,10 +256,23 @@ func NewSQLiteManager(filename, startupQueries string, config *Config) (*Manager
insertPhoneNumber: sqliteInsertPhoneNumberQuery, insertPhoneNumber: sqliteInsertPhoneNumberQuery,
deletePhoneNumber: sqliteDeletePhoneNumberQuery, deletePhoneNumber: sqliteDeletePhoneNumberQuery,
updateBilling: sqliteUpdateBillingQuery, updateBilling: sqliteUpdateBillingQuery,
},
} }
if err := initManager(manager); err != nil {
// NewSQLiteManager creates a new Manager backed by a SQLite database
func NewSQLiteManager(filename, startupQueries string, config *Config) (*Manager, error) {
parentDir := filepath.Dir(filename)
if !util.FileExists(parentDir) {
return nil, fmt.Errorf("user database directory %s does not exist or is not accessible", parentDir)
}
db, err := sql.Open("sqlite3", filename)
if err != nil {
return nil, err return nil, err
} }
return manager, nil if err := setupSQLite(db); err != nil {
return nil, err
}
if err := runSQLiteStartupQueries(db, startupQueries); err != nil {
return nil, err
}
return newManager(db, sqliteQueries, config)
} }

View File

@@ -1958,26 +1958,6 @@ func TestStoreTokenRemoveExpired(t *testing.T) {
}) })
} }
func TestStoreTokenCreatePrunesExcess(t *testing.T) {
forEachStoreBackend(t, func(t *testing.T, manager *Manager) {
require.Nil(t, manager.AddUser("phil", "mypass", RoleUser, false))
u, err := manager.User("phil")
require.Nil(t, err)
// Create several tokens
var tokenValues []string
for i := 0; i < 3; i++ {
tk, err := manager.CreateToken(u.ID, "label", time.Now().Add(time.Duration(i+1)*time.Hour), netip.MustParseAddr("1.2.3.4"), false)
require.Nil(t, err)
tokenValues = append(tokenValues, tk.Value)
}
tokens, err := manager.Tokens(u.ID)
require.Nil(t, err)
require.True(t, len(tokens) >= 3)
})
}
func TestStoreTokenUpdateLastAccess(t *testing.T) { func TestStoreTokenUpdateLastAccess(t *testing.T) {
forEachStoreBackend(t, func(t *testing.T, manager *Manager) { forEachStoreBackend(t, func(t *testing.T, manager *Manager) {
require.Nil(t, manager.AddUser("phil", "mypass", RoleUser, false)) require.Nil(t, manager.AddUser("phil", "mypass", RoleUser, false))