From 57c5a69f7d9771300d68e33eb02d1209ef1d7d4c Mon Sep 17 00:00:00 2001
From: David Dworken <david@daviddworken.com>
Date: Tue, 12 Sep 2023 18:55:13 -0700
Subject: [PATCH] Fix bug where we failed to delete pre-saved history entries
 due to race conditions causing the DB to be locked

---
 client/cmd/saveHistoryEntry.go | 24 +++++++++-------
 client/lib/lib.go              | 52 ++++++++++++++++++++--------------
 2 files changed, 44 insertions(+), 32 deletions(-)

diff --git a/client/cmd/saveHistoryEntry.go b/client/cmd/saveHistoryEntry.go
index a4aaf99..bd159d1 100644
--- a/client/cmd/saveHistoryEntry.go
+++ b/client/cmd/saveHistoryEntry.go
@@ -138,17 +138,21 @@ func saveHistoryEntry(ctx context.Context) {
 
 	// Drop any entries from pre-saving since they're no longer needed
 	if config.BetaMode {
-		tx, err := lib.MakeWhereQueryFromSearch(ctx, db, "cwd:"+entry.CurrentWorkingDirectory+" start_time:"+strconv.FormatInt(entry.StartTime.Unix(), 10))
-		if err != nil {
-			lib.CheckFatalError(fmt.Errorf("failed to query for pre-saved history entry: %w", err))
-		}
-		res := tx.Delete(&data.HistoryEntry{})
-		if res.Error != nil {
-			lib.CheckFatalError(fmt.Errorf("failed to delete pre-saved history entry (expected command=%#v): %w", entry.Command, res.Error))
-		}
-		if res.RowsAffected > 1 {
-			lib.CheckFatalError(fmt.Errorf("attempted to delete pre-saved entry, but something went wrong since we deleted %d rows", res.RowsAffected))
+		deletePresavedEntryFunc := func() error {
+			tx, err := lib.MakeWhereQueryFromSearch(ctx, db, "cwd:"+entry.CurrentWorkingDirectory+" start_time:"+strconv.FormatInt(entry.StartTime.Unix(), 10)+" end_time:1970/01/01_00:00:00_+0000")
+			if err != nil {
+				return fmt.Errorf("failed to query for pre-saved history entry: %w", err)
+			}
+			res := tx.Delete(&data.HistoryEntry{})
+			if res.Error != nil {
+				return fmt.Errorf("failed to delete pre-saved history entry (expected command=%#v): %w", entry.Command, res.Error)
+			}
+			if res.RowsAffected > 1 {
+				return fmt.Errorf("attempted to delete pre-saved entry, but something went wrong since we deleted %d rows", res.RowsAffected)
+			}
+			return nil
 		}
+		lib.CheckFatalError(lib.RetryingDbFunction(deletePresavedEntryFunc))
 	}
 
 	// Persist it locally
diff --git a/client/lib/lib.go b/client/lib/lib.go
index cf9e611..a4e8220 100644
--- a/client/lib/lib.go
+++ b/client/lib/lib.go
@@ -705,36 +705,36 @@ func normalizeEntryTimezone(entry data.HistoryEntry) data.HistoryEntry {
 	return entry
 }
 
-func ReliableDbCreate(db *gorm.DB, entry data.HistoryEntry) error {
-	entry = normalizeEntryTimezone(entry)
+func RetryingDbFunction(dbFunc func() error) error {
 	var err error = nil
 	i := 0
 	for i = 0; i < 10; i++ {
-		result := db.Create(entry)
-		err = result.Error
+		err = dbFunc()
 		if err == nil {
 			return nil
 		}
-		if err != nil {
-			errMsg := err.Error()
-			if errMsg == "database is locked (5) (SQLITE_BUSY)" || errMsg == "database is locked (261)" {
-				time.Sleep(time.Duration(i*rand.Intn(100)) * time.Millisecond)
-				continue
-			}
-			if strings.Contains(errMsg, "UNIQUE constraint failed") {
-				if i == 0 {
-					return err
-				} else {
-					return nil
-				}
-			}
-			return fmt.Errorf("unrecoverable sqlite error: %w", err)
+		errMsg := err.Error()
+		if errMsg == "database is locked (5) (SQLITE_BUSY)" || errMsg == "database is locked (261)" {
+			time.Sleep(time.Duration(i*rand.Intn(100)) * time.Millisecond)
+			continue
 		}
-		if err != nil && err.Error() != "database is locked (5) (SQLITE_BUSY)" {
-			return fmt.Errorf("unrecoverable sqlite error: %w", err)
+		if strings.Contains(errMsg, "UNIQUE constraint failed") {
+			if i == 0 {
+				return err
+			} else {
+				return nil
+			}
 		}
+		return fmt.Errorf("unrecoverable sqlite error: %w", err)
 	}
-	return fmt.Errorf("failed to create DB entry even with %d retries: %w", i, err)
+	return fmt.Errorf("failed to execute DB transaction even with %d retries: %w", i, err)
+}
+
+func ReliableDbCreate(db *gorm.DB, entry data.HistoryEntry) error {
+	entry = normalizeEntryTimezone(entry)
+	return RetryingDbFunction(func() error {
+		return db.Create(entry).Error
+	})
 }
 
 func EncryptAndMarshal(config hctx.ClientConfig, entries []*data.HistoryEntry) ([]byte, error) {
@@ -952,9 +952,17 @@ func parseAtomizedToken(ctx context.Context, token string) (string, interface{},
 		// internally for pre-saving history entries.
 		t, err := parseTimeGenerously(val)
 		if err != nil {
-			return "", nil, nil, fmt.Errorf("failed to parse after:%s as a timestamp: %w", val, err)
+			return "", nil, nil, fmt.Errorf("failed to parse start_time:%s as a timestamp: %w", val, err)
 		}
 		return "(CAST(strftime(\"%s\",start_time) AS INTEGER) = ?)", t.Unix(), nil, nil
+	case "end_time":
+		// Note that this atom probably isn't useful for interactive usage since it does exact matching, but we use it
+		// internally for pre-saving history entries.
+		t, err := parseTimeGenerously(val)
+		if err != nil {
+			return "", nil, nil, fmt.Errorf("failed to parse end_time:%s as a timestamp: %w", val, err)
+		}
+		return "(CAST(strftime(\"%s\",end_time) AS INTEGER) = ?)", t.Unix(), nil, nil
 	case "command":
 		return "(instr(command, ?) > 0)", val, nil, nil
 	default: