From f8b515c32811c553aec1a87db4bf0fc696a46801 Mon Sep 17 00:00:00 2001
From: David Dworken <david@daviddworken.com>
Date: Mon, 30 Oct 2023 17:50:47 -0700
Subject: [PATCH] Update custom column support to also automatically retry DB
 errors to further harden against issues like #119

---
 client/lib/lib.go | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/client/lib/lib.go b/client/lib/lib.go
index a0234dc..597eb00 100644
--- a/client/lib/lib.go
+++ b/client/lib/lib.go
@@ -4,6 +4,7 @@ import (
 	"bufio"
 	"bytes"
 	"context"
+	"database/sql"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -937,15 +938,17 @@ func parseAtomizedToken(ctx context.Context, token string) (string, any, any, er
 
 func getAllCustomColumnNames(ctx context.Context) ([]string, error) {
 	db := hctx.GetDb(ctx)
-	query := `
-	SELECT DISTINCT json_extract(value, '$.name') as cc_name
-	FROM history_entries 
-	JOIN json_each(custom_columns)
-	WHERE value IS NOT NULL
-	LIMIT 10`
-	rows, err := db.Raw(query).Rows()
+	rows, err := RetryingDbFunctionWithResult(func() (*sql.Rows, error) {
+		query := `
+		SELECT DISTINCT json_extract(value, '$.name') as cc_name
+		FROM history_entries 
+		JOIN json_each(custom_columns)
+		WHERE value IS NOT NULL
+		LIMIT 10`
+		return db.Raw(query).Rows()
+	})
 	if err != nil {
-		return nil, err
+		return nil, fmt.Errorf("failed to query for list of custom columns: %v", err)
 	}
 	ccNames := make([]string, 0)
 	for rows.Next() {