Add support for searching for custom columns as part of default searches (#286)

* Refactor code for searching for custom columns

* Add support for searching custom columns in default searches

* Add integration test for default searches for custom columns

* Add memoization to avoid repeated work for queries

* update comment

* Update golden

* Update goldens

client/cmd/configSet.go

@@ -281,14 +281,22 @@ var setDefaultSearchColumns = &cobra.Command{
 	Short: "Get the list of columns that are used for \"default\" search queries that don't use any search atoms",
 	Long:  "By default hishtory queries are checked against `command`, `current_working_directory`, and `hostname`. This option can be used to exclude `current_working_directory` and/or `hostname` from default search queries. E.g. `hishtory config-set default-search-columns hostname command` would exclude `current_working_directory` from default searches.",
 	Args:  cobra.OnlyValidArgs,
-	// Note: If we are ever adding new arguments to this list, we should consider adding support for this config option in configAdd.go and configDelete.go.
-	ValidArgs: []string{"current_working_directory", "hostname", "command"},
 	Run: func(cmd *cobra.Command, args []string) {
+		// TODO: Add configAdd and configDelete support for this command
 		ctx := hctx.MakeContext()
 		config := hctx.GetConf(ctx)
 		if !slices.Contains(args, "command") {
 			lib.CheckFatalError(fmt.Errorf("command is a required default search column"))
 		}
+		customColNames, err := lib.GetAllCustomColumnNames(ctx)
+		if err != nil {
+			lib.CheckFatalError(fmt.Errorf("failed to get custom column names: %v", err))
+		}
+		for _, col := range args {
+			if !slices.Contains(lib.SUPPORTED_DEFAULT_COLUMNS, col) && !slices.Contains(customColNames, col) {
+				lib.CheckFatalError(fmt.Errorf("column %q is not a valid column name", col))
+			}
+		}
 		config.DefaultSearchColumns = args
 		lib.CheckFatalError(hctx.SetConfig(config))
 	},

client/hctx/hctx.go

@@ -301,7 +301,7 @@ func GetConfig() (ClientConfig, error) {
 		config.LogLevel = logrus.InfoLevel
 	}
 	if len(config.DefaultSearchColumns) == 0 {
-		config.DefaultSearchColumns = []string{"command", "current_working_directory", "hostname"}
+		config.DefaultSearchColumns = []string{"command", "hostname", "current_working_directory"}
 	}
 	return config, nil
 }

client/integration_test.go

@@ -3551,10 +3551,15 @@ func TestDefaultSearchColumns(t *testing.T) {
 	e1 := testutils.MakeFakeHistoryEntry("echo hi")
 	e1.CurrentWorkingDirectory = "/cwd1/"
 	e1.Hostname = "h1"
+	e1.CustomColumns = make(data.CustomColumns, 2)
+	e1.CustomColumns[0] = data.CustomColumn{Name: "MyCol", Val: "baz"}
+	e1.CustomColumns[1] = data.CustomColumn{Name: "baz", Val: "bar"}
 	require.NoError(t, db.Create(e1).Error)
 	e2 := testutils.MakeFakeHistoryEntry("ls")
 	e2.CurrentWorkingDirectory = "/echo/"
 	e2.Hostname = "hi"
+	e2.CustomColumns = make(data.CustomColumns, 1)
+	e2.CustomColumns[0] = data.CustomColumn{Name: "MyCol", Val: "bar"}
 	require.NoError(t, db.Create(e2).Error)
 
 	// Check that by default all columns are included
@@ -3565,7 +3570,7 @@ func TestDefaultSearchColumns(t *testing.T) {
 
 	// Update the config value to exclude CWD
 	out = tester.RunInteractiveShell(t, ` hishtory config-get default-search-columns`)
-	require.Equal(t, out, "command current_working_directory hostname \n")
+	require.Equal(t, out, "command hostname current_working_directory \n")
 	tester.RunInteractiveShell(t, ` hishtory config-set default-search-columns 'hostname' 'command'`)
 	out = tester.RunInteractiveShell(t, ` hishtory config-get default-search-columns`)
 	require.Equal(t, out, "hostname command \n")
@@ -3586,6 +3591,23 @@ func TestDefaultSearchColumns(t *testing.T) {
 	testutils.CompareGoldens(t, out, "TestDefaultSearchColumns-NoCWDHostname-Echo")
 	out = tester.RunInteractiveShell(t, ` hishtory export hi | grep -v pipefail`)
 	testutils.CompareGoldens(t, out, "TestDefaultSearchColumns-NoCWDHostname-Hi")
+
+	// Add a custom column to the search
+	tester.RunInteractiveShell(t, ` hishtory config-set default-search-columns 'hostname' 'command' MyCol`)
+	out = tester.RunInteractiveShell(t, ` hishtory config-get default-search-columns`)
+	require.Equal(t, out, "hostname command MyCol \n")
+
+	// Check that the normal searches still work fine
+	out = tester.RunInteractiveShell(t, ` hishtory export echo | grep -v pipefail`)
+	testutils.CompareGoldens(t, out, "TestDefaultSearchColumns-NoCWD-Echo")
+	out = tester.RunInteractiveShell(t, ` hishtory export hi | grep -v pipefail`)
+	testutils.CompareGoldens(t, out, "TestDefaultSearchColumns-NoCWD-Hi")
+
+	// Check that we can search for the custom column by default
+	out = tester.RunInteractiveShell(t, ` hishtory export bar | grep -v pipefail`)
+	testutils.CompareGoldens(t, out, "TestDefaultSearchColumns-MyCol-bar")
+	out = tester.RunInteractiveShell(t, ` hishtory export baz | grep -v pipefail`)
+	testutils.CompareGoldens(t, out, "TestDefaultSearchColumns-MyCol-baz")
 }
 
 // TODO: somehow test/confirm that hishtory works even if only bash/only zsh is installed

client/lib/lib.go

@@ -784,11 +784,11 @@ func MakeWhereQueryFromSearch(ctx context.Context, db *gorm.DB, query string) (*
 				}
 				tx = where(tx, "NOT "+query, v1, v2)
 			} else {
-				query, v1, v2, v3, err := parseNonAtomizedToken(ctx, token[1:])
+				query, args, err := parseNonAtomizedToken(ctx, token[1:])
 				if err != nil {
 					return nil, err
 				}
-				tx = where(tx, "NOT "+query, v1, v2, v3)
+				tx = where(tx, "NOT "+query, args...)
 			}
 		} else if containsUnescaped(token, ":") {
 			query, v1, v2, err := parseAtomizedToken(ctx, token)
@@ -797,11 +797,11 @@ func MakeWhereQueryFromSearch(ctx context.Context, db *gorm.DB, query string) (*
 			}
 			tx = where(tx, query, v1, v2)
 		} else {
-			query, v1, v2, v3, err := parseNonAtomizedToken(ctx, token)
+			query, args, err := parseNonAtomizedToken(ctx, token)
 			if err != nil {
 				return nil, err
 			}
-			tx = where(tx, query, v1, v2, v3)
+			tx = where(tx, query, args...)
 		}
 	}
 	return tx, nil
@@ -851,36 +851,27 @@ func retryingSearch(ctx context.Context, db *gorm.DB, query string, limit, offse
 	return historyEntries, nil
 }
 
-func parseNonAtomizedToken(ctx context.Context, token string) (string, any, any, any, error) {
+var SUPPORTED_DEFAULT_COLUMNS = []string{"command", "hostname", "current_working_directory"}
+
+func parseNonAtomizedToken(ctx context.Context, token string) (string, []any, error) {
 	wildcardedToken := "%" + unescape(token) + "%"
 	query := "(false "
-	numFilters := 0
-	if slices.Contains(hctx.GetConf(ctx).DefaultSearchColumns, "command") {
-		query += "OR command LIKE ? "
-		numFilters += 1
-	}
-	if slices.Contains(hctx.GetConf(ctx).DefaultSearchColumns, "hostname") {
-		query += "OR hostname LIKE ? "
-		numFilters += 1
-	}
-	if slices.Contains(hctx.GetConf(ctx).DefaultSearchColumns, "current_working_directory") {
-		query += "OR current_working_directory LIKE ? "
-		numFilters += 1
+	args := make([]any, 0)
+	for _, column := range hctx.GetConf(ctx).DefaultSearchColumns {
+		if slices.Contains(SUPPORTED_DEFAULT_COLUMNS, column) {
+			query += "OR " + column + " LIKE ? "
+			args = append(args, wildcardedToken)
+		} else {
+			q, a, err := buildCustomColumnSearchQuery(ctx, column, unescape(token))
+			if err != nil {
+				return "", nil, err
+			}
+			query += "OR " + q + " "
+			args = append(args, a...)
+		}
 	}
 	query += ")"
-	var t1 any = nil
-	var t2 any = nil
-	var t3 any = nil
-	if numFilters >= 1 {
-		t1 = wildcardedToken
-	}
-	if numFilters >= 2 {
-		t2 = wildcardedToken
-	}
-	if numFilters >= 3 {
-		t3 = wildcardedToken
-	}
-	return query, t1, t2, t3, nil
+	return query, args, nil
 }
 
 func parseAtomizedToken(ctx context.Context, token string) (string, any, any, error) {
@@ -932,34 +923,54 @@ func parseAtomizedToken(ctx context.Context, token string) (string, any, any, er
 	case "command":
 		return "(instr(command, ?) > 0)", val, nil, nil
 	default:
-		knownCustomColumns := make([]string, 0)
-		// Get custom columns that are defined on this machine
-		conf := hctx.GetConf(ctx)
-		for _, c := range conf.CustomColumns {
-			knownCustomColumns = append(knownCustomColumns, c.ColumnName)
-		}
-		// Also get all ones that are in the DB
-		names, err := getAllCustomColumnNames(ctx)
+		q, args, err := buildCustomColumnSearchQuery(ctx, field, val)
 		if err != nil {
-			return "", nil, nil, fmt.Errorf("failed to get custom column names from the DB: %w", err)
+			return "", nil, nil, err
 		}
-		knownCustomColumns = append(knownCustomColumns, names...)
-		// Check if the atom is for a custom column that exists and if it isn't, return an error
-		isCustomColumn := false
-		for _, ccName := range knownCustomColumns {
-			if ccName == field {
-				isCustomColumn = true
-			}
+		if len(args) != 2 {
+			return "", nil, nil, fmt.Errorf("custom column search query returned an unexpected number of args: %d", len(args))
 		}
-		if !isCustomColumn {
-			return "", nil, nil, fmt.Errorf("search query contains unknown search atom '%s' that doesn't match any column names", field)
-		}
-		// Build the where clause for the custom column
-		return "EXISTS (SELECT 1 FROM json_each(custom_columns) WHERE json_extract(value, '$.name') = ? and instr(json_extract(value, '$.value'), ?) > 0)", field, val, nil
+		return q, args[0], args[1], nil
 	}
 }
 
-func getAllCustomColumnNames(ctx context.Context) ([]string, error) {
+func buildCustomColumnSearchQuery(ctx context.Context, columnName, columnVal string) (string, []any, error) {
+	knownCustomColumns, err := GetAllCustomColumnNames(ctx)
+	if err != nil {
+		return "", nil, fmt.Errorf("failed to get list of known custom columns: %w", err)
+	}
+	if !slices.Contains(knownCustomColumns, columnName) {
+		return "", nil, fmt.Errorf("search query contains unknown search atom '%s' that doesn't match any column names", columnName)
+	}
+	// Build the where clause for the custom column
+	return "EXISTS (SELECT 1 FROM json_each(custom_columns) WHERE json_extract(value, '$.name') = ? and instr(json_extract(value, '$.value'), ?) > 0)", []any{columnName, columnVal}, nil
+}
+
+func GetAllCustomColumnNames(ctx context.Context) ([]string, error) {
+	knownCustomColumns := make([]string, 0)
+	// Get custom columns that are defined on this machine
+	conf := hctx.GetConf(ctx)
+	for _, c := range conf.CustomColumns {
+		knownCustomColumns = append(knownCustomColumns, c.ColumnName)
+	}
+	// Also get all ones that are in the DB
+	names, err := getAllCustomColumnNamesFromDb(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get custom column names from the DB: %w", err)
+	}
+	knownCustomColumns = append(knownCustomColumns, names...)
+	return knownCustomColumns, nil
+}
+
+var cachedCustomColumnNames []string
+
+func getAllCustomColumnNamesFromDb(ctx context.Context) ([]string, error) {
+	if len(cachedCustomColumnNames) > 0 {
+		// Note: We memoize this function since it is called repeatedly in the TUI and querying the
+		// entire DB for every updated search is quite inefficient. This is reasonable since the set
+		// of custom columns shouldn't ever change within the lifetime of one hishtory process.
+		return cachedCustomColumnNames, nil
+	}
 	db := hctx.GetDb(ctx)
 	rows, err := RetryingDbFunctionWithResult(func() (*sql.Rows, error) {
 		query := `
@@ -981,6 +992,7 @@ func getAllCustomColumnNames(ctx context.Context) ([]string, error) {
 		}
 		ccNames = append(ccNames, ccName)
 	}
+	cachedCustomColumnNames = ccNames
 	return ccNames, nil
 }
 

client/lib/lib_test.go

@@ -353,31 +353,31 @@ func TestParseNonAtomizedToken(t *testing.T) {
 	ctx := hctx.MakeContext()
 
 	// Default
-	q, v1, v2, v3, err := parseNonAtomizedToken(ctx, "echo hello")
+	q, args, err := parseNonAtomizedToken(ctx, "echo hello")
 	require.NoError(t, err)
 	require.Equal(t, "(false OR command LIKE ? OR hostname LIKE ? OR current_working_directory LIKE ? )", q)
-	require.Equal(t, v1, "%echo hello%")
-	require.Equal(t, v2, "%echo hello%")
-	require.Equal(t, v3, "%echo hello%")
+	require.Len(t, args, 3)
+	require.Equal(t, args[0], "%echo hello%")
+	require.Equal(t, args[1], "%echo hello%")
+	require.Equal(t, args[2], "%echo hello%")
 
 	// Skipping cwd
 	config := hctx.GetConf(ctx)
-	config.DefaultSearchColumns = []string{"hostname", "command"}
-	q, v1, v2, v3, err = parseNonAtomizedToken(ctx, "echo hello")
+	config.DefaultSearchColumns = []string{"command", "hostname"}
+	q, args, err = parseNonAtomizedToken(ctx, "echo hello")
 	require.NoError(t, err)
 	require.Equal(t, "(false OR command LIKE ? OR hostname LIKE ? )", q)
-	require.Equal(t, v1, "%echo hello%")
-	require.Equal(t, v2, "%echo hello%")
-	require.Nil(t, v3)
+	require.Len(t, args, 2)
+	require.Equal(t, args[0], "%echo hello%")
+	require.Equal(t, args[1], "%echo hello%")
 
 	// Skipping cwd and hostname
 	config.DefaultSearchColumns = []string{"command"}
-	q, v1, v2, v3, err = parseNonAtomizedToken(ctx, "echo hello")
+	q, args, err = parseNonAtomizedToken(ctx, "echo hello")
 	require.NoError(t, err)
 	require.Equal(t, "(false OR command LIKE ? )", q)
-	require.Equal(t, v1, "%echo hello%")
-	require.Nil(t, v2)
-	require.Nil(t, v3)
+	require.Len(t, args, 1)
+	require.Equal(t, args[0], "%echo hello%")
 }
 
 func TestWhere(t *testing.T) {

client/testdata/TestDefaultSearchColumns-MyCol-bar

@@ -0,0 +1 @@
+ls

client/testdata/TestDefaultSearchColumns-MyCol-baz

@@ -0,0 +1 @@
+echo hi

client/testdata/TestStatusFullConfig

@@ -70,6 +70,6 @@ Full Config:
 	fullscreenrendering: false
 	defaultsearchcolumns:
 	    - command
-	    - current_working_directory
 	    - hostname
+	    - current_working_directory