feat: Delete All User Data (RTBF) (#225)

* delete all user data

* check return from rollback

Author: danielchalef

pkg/store/postgres/purge_test.go

@@ -7,7 +7,7 @@ import (
 )
 
 func TestPurgeDeleted(t *testing.T) {
-	sessionID, err := setupTestDeleteData(testCtx, testDB)
+	sessionID, err := setupSessionDeleteTestData(testCtx, testDB, "")
 	assert.NoError(t, err, "setupTestDeleteData should not return an error")
 
 	sessionStore := NewSessionDAO(testDB)

pkg/store/postgres/session.go

@@ -201,41 +201,66 @@ func (dao *SessionDAO) updateSession(
 }
 
 // Delete soft-deletes a session from the database by its sessionID.
-// It also soft-deletes all messages and message embeddings associated with the session.
+// It also soft-deletes all messages, message embeddings, and summaries associated with the session.
 func (dao *SessionDAO) Delete(ctx context.Context, sessionID string) error {
 	dbSession := &SessionSchema{}
 
-	r, err := dao.db.NewDelete().
+	tx, err := dao.db.BeginTx(ctx, nil)
+	if err != nil {
+		return fmt.Errorf("failed to begin transaction: %w", err)
+	}
+
+	r, err := tx.NewDelete().
 		Model(dbSession).
 		Where("session_id = ?", sessionID).
 		Exec(ctx)
 	if err != nil {
+		rollbackErr := tx.Rollback()
+		if rollbackErr != nil {
+			return fmt.Errorf("failed to delete session: %v, failed to rollback transaction: %w", err, rollbackErr)
+		}
 		return fmt.Errorf("failed to delete session: %w", err)
 	}
 
 	rowsAffected, err := r.RowsAffected()
 	if err != nil {
+		rollbackErr := tx.Rollback()
+		if rollbackErr != nil {
+			return fmt.Errorf("failed to delete session: %v, failed to rollback transaction: %w", err, rollbackErr)
+		}
 		return fmt.Errorf("failed to get rows affected: %w", err)
 	}
 	if rowsAffected == 0 {
+		rollbackErr := tx.Rollback()
+		if rollbackErr != nil {
+			return fmt.Errorf("failed to delete session: %v, failed to rollback transaction: %w", err, rollbackErr)
+		}
 		return models.NewNotFoundError("session " + sessionID)
 	}
 
-	// delete all messages and message embeddings associated with the session
+	// delete all messages, message embeddings, and summaries associated with the session
 	for _, schema := range messageTableList {
 		if _, ok := schema.(*SessionSchema); ok {
 			continue
 		}
 		log.Debugf("deleting session %s from schema %T", sessionID, schema)
-		_, err := dao.db.NewDelete().
+		_, err := tx.NewDelete().
 			Model(schema).
 			Where("session_id = ?", sessionID).
 			Exec(ctx)
 		if err != nil {
+			rollbackErr := tx.Rollback()
+			if rollbackErr != nil {
+				return fmt.Errorf("failed to delete session: %v, failed to rollback transaction: %w", err, rollbackErr)
+			}
 			return fmt.Errorf("error deleting rows from %T: %w", schema, err)
 		}
 	}
 
+	if err := tx.Commit(); err != nil {
+		return fmt.Errorf("failed to commit transaction: %w", err)
+	}
+
 	return nil
 }
 

pkg/store/postgres/session_test.go

@@ -246,7 +246,7 @@ func TestSessionDAO_DeleteSessionDeletesSummaryMessages(t *testing.T) {
 
 	sessionStore := NewSessionDAO(testDB)
 
-	sessionID, err := setupTestDeleteData(testCtx, testDB)
+	sessionID, err := setupSessionDeleteTestData(testCtx, testDB, "")
 	assert.NoError(t, err, "setupTestDeleteData should not return an error")
 
 	err = sessionStore.Delete(testCtx, sessionID)
@@ -268,7 +268,7 @@ func TestSessionDAO_DeleteSessionDeletesSummaryMessages(t *testing.T) {
 }
 
 func TestSessionDAO_UndeleteSession(t *testing.T) {
-	sessionID, err := setupTestDeleteData(testCtx, testDB)
+	sessionID, err := setupSessionDeleteTestData(testCtx, testDB, "")
 	assert.NoError(t, err, "setupTestDeleteData should not return an error")
 
 	sessionStore := NewSessionDAO(testDB)
@@ -292,16 +292,22 @@ func TestSessionDAO_UndeleteSession(t *testing.T) {
 	assert.Nil(t, respMessages, "getMessages should return nil")
 }
 
-func setupTestDeleteData(ctx context.Context, testDB *bun.DB) (string, error) {
+func setupSessionDeleteTestData(ctx context.Context, testDB *bun.DB, userID string) (string, error) {
 	// Test data
 	sessionID, err := testutils.GenerateRandomSessionID(16)
 	if err != nil {
 		return "", err
 	}
 
+	var userIDPtr *string
+	if userID != "" {
+		userIDPtr = &userID
+	}
+
 	dao := NewSessionDAO(testDB)
 	_, err = dao.Create(ctx, &models.CreateSessionRequest{
 		SessionID: sessionID,
+		UserID:    userIDPtr,
 	})
 	if err != nil {
 		return "", err

pkg/store/postgres/userstore.go

@@ -170,18 +170,65 @@ func (dao *UserStoreDAO) updateUser(
 
 // Delete deletes a user.
 func (dao *UserStoreDAO) Delete(ctx context.Context, userID string) error {
+	// Start a new transaction
+	tx, err := dao.db.Begin()
+	if err != nil {
+		return err
+	}
+
+	// Delete all related sessions
+	sessions, err := dao.GetSessions(ctx, userID)
+	if err != nil {
+		rollbackErr := tx.Rollback()
+		if rollbackErr != nil {
+			return fmt.Errorf("failed to delete user: %v, failed to rollback transaction: %w", err, rollbackErr)
+		}
+		return err
+	}
+
+	sessionStore := NewSessionDAO(dao.db)
+	for s := range sessions {
+		err := sessionStore.Delete(ctx, sessions[s].SessionID)
+		if err != nil {
+			rollbackErr := tx.Rollback()
+			if rollbackErr != nil {
+				return fmt.Errorf("failed to delete user: %v, failed to rollback transaction: %w", err, rollbackErr)
+			}
+			return err
+		}
+	}
+
+	// Delete User
 	r, err := dao.db.NewDelete().Model(&models.User{}).Where("user_id = ?", userID).Exec(ctx)
 	if err != nil {
+		rollbackErr := tx.Rollback()
+		if rollbackErr != nil {
+			return fmt.Errorf("failed to delete user: %v, failed to rollback transaction: %w", err, rollbackErr)
+		}
 		return err
 	}
 	rowsAffected, err := r.RowsAffected()
 	if err != nil {
+		rollbackErr := tx.Rollback()
+		if rollbackErr != nil {
+			return fmt.Errorf("failed to delete user: %v, failed to rollback transaction: %w", err, rollbackErr)
+		}
 		return err
 	}
 	if rowsAffected == 0 {
+		rollbackErr := tx.Rollback()
+		if rollbackErr != nil {
+			return fmt.Errorf("failed to delete user: %v, failed to rollback transaction: %w", err, rollbackErr)
+		}
 		return models.NewNotFoundError("user " + userID)
 	}
 
+	// Commit the transaction
+	err = tx.Commit()
+	if err != nil {
+		return err
+	}
+
 	return nil
 }
 

pkg/store/postgres/userstore_test.go

@@ -135,11 +135,40 @@ func TestUserStoreDAO(t *testing.T) {
 
 	// Test Delete
 	t.Run("Delete", func(t *testing.T) {
+		testSessions := []string{}
+		for i := 0; i < 2; i++ {
+			sessionID, err := setupSessionDeleteTestData(testCtx, testDB, user.UserID)
+			assert.NoError(t, err, "setupTestDeleteData should not return an error")
+			testSessions = append(testSessions, sessionID)
+		}
+
 		err := userStore.Delete(ctx, user.UserID)
 		assert.NoError(t, err)
 
 		_, err = userStore.Get(ctx, user.UserID)
 		assert.ErrorIs(t, err, models.ErrNotFound)
+
+		// Check that all related sessions are deleted
+		retSessions, err := userStore.GetSessions(ctx, user.UserID)
+		assert.NoError(t, err)
+		assert.Equal(t, 0, len(retSessions))
+
+		// Test that messages and summaries are deleted
+		for _, sessionID := range testSessions {
+			respMessages, err := getMessages(testCtx, testDB, sessionID, 999, nil, 999)
+			assert.NoError(t, err, "getMessages should not return an error")
+			assert.Nil(t, respMessages, "getMessages should return nil")
+
+			// Test that summary is deleted
+			respSummary, err := getSummary(testCtx, testDB, sessionID)
+			assert.NoError(t, err, "getSummary should not return an error")
+			assert.Nil(t, respSummary, "getSummary should return nil")
+
+			// check that embeddings are deleted
+			respEmbeddings, err := getMessageEmbeddings(testCtx, testDB, sessionID)
+			assert.NoError(t, err, "getMessageEmbeddings should not return an error")
+			assert.Equal(t, 0, len(respEmbeddings), "getMessageEmbeddings should return 0 results")
+		}
 	})
 
 	t.Run("Delete Non-Existant Session should result in NotFoundError", func(t *testing.T) {