From ea541c126219e2f11ae9a97d39a33555a4583397 Mon Sep 17 00:00:00 2001 From: Umputun Date: Fri, 10 Jan 2025 18:38:22 -0600 Subject: [PATCH] Refactor detected spam storage to support multiple database types - Replace single query string with a query map for different DB engines - Use `engine.PickQuery` to select appropriate queries based on DB type - Refactor schema and index creation logic for better error handling - Use `ds.db.Adopt` for query adoption, enhancing flexibility - Add a check for nil DB connection during DetectedSpam initialization --- app/storage/detected_spam.go | 88 +++++++++++++++++++++++-------- app/storage/detected_spam_test.go | 10 ++-- 2 files changed, 72 insertions(+), 26 deletions(-) diff --git a/app/storage/detected_spam.go b/app/storage/detected_spam.go index ed3c251..8b1a73a 100644 --- a/app/storage/detected_spam.go +++ b/app/storage/detected_spam.go @@ -34,30 +34,70 @@ type DetectedSpamInfo struct { Timestamp time.Time `db:"timestamp"` Added bool `db:"added"` // added to samples ChecksJSON string `db:"checks"` // Store as JSON - Checks []spamcheck.Response `db:"-"` // Don't store in DB directly, for db it uses ChecksJSON + Checks []spamcheck.Response `db:"-"` // Don't store in DB directly } -var detectedSpamSchema = ` - CREATE TABLE IF NOT EXISTS detected_spam ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - gid TEXT NOT NULL DEFAULT '', - text TEXT, - user_id INTEGER, - user_name TEXT, - timestamp DATETIME DEFAULT CURRENT_TIMESTAMP, - added BOOLEAN DEFAULT 0, - checks TEXT - ); - CREATE INDEX IF NOT EXISTS idx_detected_spam_timestamp ON detected_spam(timestamp); - CREATE INDEX IF NOT EXISTS idx_detected_spam_gid ON detected_spam(gid) -` +// CmdCreateDetectedSpamTable creates detected_spam table +const CmdCreateDetectedSpamTable engine.DBCmd = iota + 200 + +// queries holds all detected spam queries +var detectedSpamQueries = engine.QueryMap{ + engine.Sqlite: { + CmdCreateDetectedSpamTable: ` + CREATE TABLE IF NOT EXISTS detected_spam ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + gid TEXT NOT NULL DEFAULT '', + text TEXT, + user_id INTEGER, + user_name TEXT, + timestamp DATETIME DEFAULT CURRENT_TIMESTAMP, + added BOOLEAN DEFAULT 0, + checks TEXT + )`, + CmdAddGIDColumn: "ALTER TABLE detected_spam ADD COLUMN gid TEXT DEFAULT ''", + }, + engine.Postgres: { + CmdCreateDetectedSpamTable: ` + CREATE TABLE IF NOT EXISTS detected_spam ( + id SERIAL PRIMARY KEY, + gid TEXT NOT NULL DEFAULT '', + text TEXT, + user_id BIGINT, + user_name TEXT, + timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP, + added BOOLEAN DEFAULT false, + checks TEXT + )`, + CmdAddGIDColumn: "ALTER TABLE detected_spam ADD COLUMN IF NOT EXISTS gid TEXT DEFAULT ''", + }, +} // NewDetectedSpam creates a new DetectedSpam storage func NewDetectedSpam(ctx context.Context, db *engine.SQL) (*DetectedSpam, error) { - err := engine.InitDB(ctx, db, "detected_spam", detectedSpamSchema, migrateDetectedSpamTx) + if db == nil { + return nil, fmt.Errorf("db connection is nil") + } + + // get schema creation queries + schema, err := engine.PickQuery(detectedSpamQueries, db.Type(), CmdCreateDetectedSpamTable) + if err != nil { + return nil, fmt.Errorf("failed to get create table query: %w", err) + } + + // initialize db with schema + err = engine.InitDB(ctx, db, "detected_spam", schema, migrateDetectedSpamTx) if err != nil { return nil, err } + + // create indexes after migration when all columns exist + createIndexes := ` + CREATE INDEX IF NOT EXISTS idx_detected_spam_timestamp ON detected_spam(timestamp); + CREATE INDEX IF NOT EXISTS idx_detected_spam_gid ON detected_spam(gid)` + if _, err = db.ExecContext(ctx, createIndexes); err != nil { + return nil, fmt.Errorf("failed to create indexes: %w", err) + } + return &DetectedSpam{db: db, RWLocker: db.MakeLock()}, nil } @@ -75,8 +115,9 @@ func (ds *DetectedSpam) Write(ctx context.Context, entry DetectedSpamInfo, check return fmt.Errorf("failed to marshal checks: %w", err) } - query := `INSERT INTO detected_spam (gid, text, user_id, user_name, timestamp, checks) VALUES (?, ?, ?, ?, ?, ?)` - if _, err := ds.db.ExecContext(ctx, query, entry.GID, entry.Text, entry.UserID, entry.UserName, entry.Timestamp, checksJSON); err != nil { + query := ds.db.Adopt("INSERT INTO detected_spam (gid, text, user_id, user_name, timestamp, checks) VALUES (?, ?, ?, ?, ?, ?)") + _, err = ds.db.ExecContext(ctx, query, entry.GID, entry.Text, entry.UserID, entry.UserName, entry.Timestamp, string(checksJSON)) + if err != nil { return fmt.Errorf("failed to insert detected spam entry: %w", err) } @@ -89,8 +130,8 @@ func (ds *DetectedSpam) SetAddedToSamplesFlag(ctx context.Context, id int64) err ds.Lock() defer ds.Unlock() - query := `UPDATE detected_spam SET added = 1 WHERE id = ?` - if _, err := ds.db.ExecContext(ctx, query, id); err != nil { + query := ds.db.Adopt("UPDATE detected_spam SET added = ? WHERE id = ?") + if _, err := ds.db.ExecContext(ctx, query, true, id); err != nil { return fmt.Errorf("failed to update added to samples flag: %w", err) } return nil @@ -101,8 +142,9 @@ func (ds *DetectedSpam) Read(ctx context.Context) ([]DetectedSpamInfo, error) { ds.RLock() defer ds.RUnlock() + query := ds.db.Adopt("SELECT * FROM detected_spam ORDER BY timestamp DESC LIMIT ?") var entries []DetectedSpamInfo - err := ds.db.SelectContext(ctx, &entries, "SELECT * FROM detected_spam ORDER BY timestamp DESC LIMIT ?", maxDetectedSpamEntries) + err := ds.db.SelectContext(ctx, &entries, query, maxDetectedSpamEntries) if err != nil { return nil, fmt.Errorf("failed to get detected spam entries: %w", err) } @@ -123,9 +165,9 @@ func (ds *DetectedSpam) FindByUserID(ctx context.Context, userID int64) (*Detect ds.RLock() defer ds.RUnlock() + query := ds.db.Adopt("SELECT * FROM detected_spam WHERE user_id = ? AND gid = ? ORDER BY timestamp DESC LIMIT 1") var entry DetectedSpamInfo - err := ds.db.GetContext(ctx, &entry, "SELECT * FROM detected_spam WHERE user_id = ? AND gid = ?"+ - " ORDER BY timestamp DESC LIMIT 1", userID, ds.db.GID()) + err := ds.db.GetContext(ctx, &entry, query, userID, ds.db.GID()) if errors.Is(err, sql.ErrNoRows) { // not found, return nil *DetectedSpamInfo instead of error return nil, nil diff --git a/app/storage/detected_spam_test.go b/app/storage/detected_spam_test.go index a8f9da9..59913eb 100644 --- a/app/storage/detected_spam_test.go +++ b/app/storage/detected_spam_test.go @@ -768,14 +768,18 @@ func TestDetectedSpam_MigrationEdgeCases(t *testing.T) { db, teardown := setupTestDB(t) defer teardown() + // get schema creation queries + schema, err := engine.PickQuery(detectedSpamQueries, db.Type(), CmdCreateDetectedSpamTable) + require.NoError(t, err) + // create schema with gid already present - _, err := db.Exec(detectedSpamSchema) + _, err = db.Exec(schema) require.NoError(t, err) // insert test data with existing gid _, err = db.Exec(` - INSERT INTO detected_spam (text, user_id, user_name, gid, checks) - VALUES (?, ?, ?, ?, ?)`, + INSERT INTO detected_spam (text, user_id, user_name, gid, checks) + VALUES (?, ?, ?, ?, ?)`, "spam", 1, "user1", "existing_gid", `[{"Name":"test","Spam":true}]`) require.NoError(t, err)