feat: Add regex search to task logs API (#9994)

determined-ai · Oct 4, 2024 · f5400ea · f5400ea
1 parent ddca766
commit f5400ea
Show file tree

Hide file tree

Showing 16 changed files with 971 additions and 840 deletions.
diff --git a/harness/determined/common/api/bindings.py b/harness/determined/common/api/bindings.py
diff --git a/master/internal/api/filters.go b/master/internal/api/filters.go
@@ -14,6 +14,8 @@ const (
 	FilterOperationLessThanEqual
 	// FilterOperationStringContainment checks if the field contains a value as a substring.
 	FilterOperationStringContainment
+	// FilterOperationRegexContainment checks if the field contains the regex.
+	FilterOperationRegexContainment
 )
 
 // Filter is a general representation for a filter provided to an API.

diff --git a/master/internal/api_tasks.go b/master/internal/api_tasks.go
@@ -729,11 +729,19 @@ func constructTaskLogsFilters(req *apiv1.TaskLogsRequest) ([]api.Filter, error)
 	}
 
 	if req.SearchText != "" {
-		filters = append(filters, api.Filter{
-			Field:     "log",
-			Operation: api.FilterOperationStringContainment,
-			Values:    req.SearchText,
-		})
+		if req.EnableRegex {
+			filters = append(filters, api.Filter{
+				Field:     "log",
+				Operation: api.FilterOperationRegexContainment,
+				Values:    req.SearchText,
+			})
+		} else {
+			filters = append(filters, api.Filter{
+				Field:     "log",
+				Operation: api.FilterOperationStringContainment,
+				Values:    req.SearchText,
+			})
+		}
 	}
 	return filters, nil
 }

diff --git a/master/internal/api_tasks_intg_test.go b/master/internal/api_tasks_intg_test.go
@@ -103,6 +103,23 @@ func TestPostTaskLogs(t *testing.T) {
 		require.Equal(t, e.Source, a.Source)
 		require.Equal(t, e.Stdtype, a.Stdtype)
 	}
+
+	// Test log filtering by regex
+	stream = &mockStream[*apiv1.TaskLogsResponse]{ctx: ctx}
+	err = api.TaskLogs(&apiv1.TaskLogsRequest{
+		TaskId:     string(task.TaskID),
+		SearchText: "^lo.{4}xt",
+	}, stream)
+	require.NoError(t, err)
+	require.Empty(t, stream.getData())
+
+	err = api.TaskLogs(&apiv1.TaskLogsRequest{
+		TaskId:      string(task.TaskID),
+		SearchText:  "^lo.{4}xt",
+		EnableRegex: true,
+	}, stream)
+	require.NoError(t, err)
+	require.Len(t, stream.getData(), 1)
 }
 
 func mockNotebookWithWorkspaceID(

diff --git a/master/internal/api_trials.go b/master/internal/api_trials.go
@@ -204,6 +204,7 @@ func (a *apiServer) TrialLogs(
 				TimestampAfter:  req.TimestampAfter,
 				OrderBy:         req.OrderBy,
 				SearchText:      req.SearchText,
+				EnableRegex:     req.EnableRegex,
 			}, res)
 			err := processBatches(res, func(b api.Batch) error {
 				return b.ForEach(func(i interface{}) error {

diff --git a/master/internal/db/postgres_filters.go b/master/internal/db/postgres_filters.go
@@ -75,6 +75,8 @@ func filterToSQL(
 		return fmt.Sprintf("AND encode(%s::bytea, 'escape') ILIKE  ('%%%%' || $%d || '%%%%')",
 			field,
 			paramID)
+	case api.FilterOperationRegexContainment:
+		return fmt.Sprintf("AND encode(%s::bytea, 'escape') ~ $%d", field, paramID)
 	default:
 		panic(fmt.Sprintf("cannot convert operation %d to SQL", f.Operation))
 	}

diff --git a/master/internal/db/postgres_tasks_intg_test.go b/master/internal/db/postgres_tasks_intg_test.go
@@ -733,6 +733,31 @@ func TestTaskLogsFlow(t *testing.T) {
 	require.NoError(t, err)
 	require.Len(t, logs, 2)
 
+	// Filter by search text.
+	logs, _, err = db.TaskLogs(t1In.TaskID, 5, []api.Filter{{
+		Field:     "log",
+		Operation: api.FilterOperationStringContainment,
+		Values:    []string{"this"},
+	}}, apiv1.OrderBy_ORDER_BY_UNSPECIFIED, nil)
+	require.NoError(t, err)
+	require.Len(t, logs, 2)
+
+	logs, _, err = db.TaskLogs(t1In.TaskID, 5, []api.Filter{{
+		Field:     "log",
+		Operation: api.FilterOperationStringContainment,
+		Values:    []string{"^th.s"},
+	}}, apiv1.OrderBy_ORDER_BY_UNSPECIFIED, nil)
+	require.NoError(t, err)
+	require.Empty(t, logs)
+
+	logs, _, err = db.TaskLogs(t1In.TaskID, 5, []api.Filter{{
+		Field:     "log",
+		Operation: api.FilterOperationRegexContainment,
+		Values:    []string{"^th.s"},
+	}}, apiv1.OrderBy_ORDER_BY_UNSPECIFIED, nil)
+	require.NoError(t, err)
+	require.Len(t, logs, 2)
+
 	// Test DeleteTaskLogs.
 	err = db.DeleteTaskLogs([]model.TaskID{t2In.TaskID})
 	require.NoError(t, err)

diff --git a/master/internal/elastic/elastic_task_logs.go b/master/internal/elastic/elastic_task_logs.go
@@ -463,7 +463,15 @@ func filtersToElastic(fs []api.Filter) []jsonObj {
 						},
 					},
 				})
-
+		case api.FilterOperationRegexContainment:
+			terms = append(terms,
+				jsonObj{
+					"regexp": jsonObj{
+						f.Field: jsonObj{
+							"value": fmt.Sprintf("%s", f.Values),
+						},
+					},
+				})
 		default:
 			panic(fmt.Sprintf("unsupported filter operation: %d", f.Operation))
 		}