Support where_pending and where_errored filters in GetModelSplits

cli/cmd/project/splits.go

@@ -10,7 +10,7 @@ import (
 
 func SplitsCmd(ch *cmdutil.Helper) *cobra.Command {
 	var project, path, model string
-	var local bool
+	var pending, errored, local bool
 	var pageSize uint32
 	var pageToken string
 
@@ -42,6 +42,8 @@ func SplitsCmd(ch *cmdutil.Helper) *cobra.Command {
 			res, err := rt.GetModelSplits(cmd.Context(), &runtimev1.GetModelSplitsRequest{
 				InstanceId: instanceID,
 				Model:      model,
+				Pending:    pending,
+				Errored:    errored,
 				PageSize:   pageSize,
 				PageToken:  pageToken,
 			})
@@ -64,6 +66,9 @@ func SplitsCmd(ch *cmdutil.Helper) *cobra.Command {
 	splitsCmd.Flags().StringVar(&project, "project", "", "Project Name")
 	splitsCmd.Flags().StringVar(&path, "path", ".", "Project directory")
 	splitsCmd.Flags().StringVar(&model, "model", "", "Model Name")
+	splitsCmd.Flags().BoolVar(&pending, "pending", false, "Only fetch pending splits")
+	splitsCmd.Flags().BoolVar(&errored, "errored", false, "Only fetch errored splits")
+	splitsCmd.MarkFlagsMutuallyExclusive("pending", "errored")
 	splitsCmd.Flags().BoolVar(&local, "local", false, "Target locally running Rill")
 	splitsCmd.Flags().Uint32Var(&pageSize, "page-size", 50, "Number of splits to return per page")
 	splitsCmd.Flags().StringVar(&pageToken, "page-token", "", "Pagination token")

proto/gen/rill/runtime/v1/runtime.swagger.yaml

@@ -993,6 +993,14 @@ paths:
           in: path
           required: true
           type: string
+        - name: pending
+          in: query
+          required: false
+          type: boolean
+        - name: errored
+          in: query
+          required: false
+          type: boolean
         - name: pageSize
           in: query
           required: false

proto/rill/runtime/v1/api.proto

@@ -664,6 +664,8 @@ message GetExploreResponse {
 message GetModelSplitsRequest {
   string instance_id = 1;
   string model = 2;
+  bool pending = 5;
+  bool errored = 6;
   uint32 page_size = 3 [(validate.rules).uint32 = {ignore_empty: true, lte: 10000}];
   string page_token = 4;
 }

runtime/drivers/catalog.go

@@ -39,9 +39,8 @@ type CatalogStore interface {
 	DeleteResource(ctx context.Context, v int64, k, n string) error
 	DeleteResources(ctx context.Context) error
 
-	FindModelSplits(ctx context.Context, modelID string, afterIndex int, afterKey string, limit int) ([]ModelSplit, error)
+	FindModelSplits(ctx context.Context, opts *FindModelSplitsOptions) ([]ModelSplit, error)
 	FindModelSplitsByKeys(ctx context.Context, modelID string, keys []string) ([]ModelSplit, error)
-	FindModelSplitsByPending(ctx context.Context, modelID string, limit int) ([]ModelSplit, error)
 	CheckModelSplitsHaveErrors(ctx context.Context, modelID string) (bool, error)
 	InsertModelSplit(ctx context.Context, modelID string, split ModelSplit) error
 	UpdateModelSplit(ctx context.Context, modelID string, split ModelSplit) error
@@ -78,3 +77,13 @@ type ModelSplit struct {
 	// Elapsed is the duration of the last execution of the split.
 	Elapsed time.Duration
 }
+
+// FindModelSplitsOptions is used to filter model splits.
+type FindModelSplitsOptions struct {
+	ModelID      string
+	Limit        int
+	WherePending bool
+	WhereErrored bool
+	AfterIndex   int
+	AfterKey     string
+}

runtime/drivers/catalog_test.go

@@ -39,7 +39,7 @@ func testCatalogSplits(t *testing.T, catalog drivers.CatalogStore) {
 	err = catalog.InsertModelSplit(ctx, modelID, split)
 	require.NoError(t, err)
 
-	splits, err = catalog.FindModelSplitsByPending(ctx, modelID, 10)
+	splits, err = catalog.FindModelSplits(ctx, &drivers.FindModelSplitsOptions{ModelID: modelID, WherePending: true, Limit: 10})
 	require.NoError(t, err)
 	require.Len(t, splits, 1)
 	requireSplitEqual(t, split, splits[0])

runtime/drivers/duckdb/catalogv2.go

@@ -189,18 +189,14 @@ func (c *connection) DeleteResources(ctx context.Context) error {
 	return nil
 }
 
-func (c *connection) FindModelSplits(ctx context.Context, modelID string, afterIndex int, afterKey string, limit int) ([]drivers.ModelSplit, error) {
+func (c *connection) FindModelSplits(ctx context.Context, opts *drivers.FindModelSplitsOptions) ([]drivers.ModelSplit, error) {
 	return nil, drivers.ErrNotImplemented
 }
 
 func (c *connection) FindModelSplitsByKeys(ctx context.Context, modelID string, keys []string) ([]drivers.ModelSplit, error) {
 	return nil, drivers.ErrNotImplemented
 }
 
-func (c *connection) FindModelSplitsByPending(ctx context.Context, modelID string, limit int) ([]drivers.ModelSplit, error) {
-	return nil, drivers.ErrNotImplemented
-}
-
 func (c *connection) CheckModelSplitsHaveErrors(ctx context.Context, modelID string) (bool, error) {
 	return false, drivers.ErrNotImplemented
 }

runtime/drivers/sqlite/catalog.go

@@ -142,17 +142,34 @@ func (c *catalogStore) DeleteResources(ctx context.Context) error {
 	return nil
 }
 
-func (c *catalogStore) FindModelSplits(ctx context.Context, modelID string, afterIndex int, afterKey string, limit int) ([]drivers.ModelSplit, error) {
-	rows, err := c.db.QueryContext(
-		ctx,
-		"SELECT key, data_json, idx, watermark, executed_on, error, elapsed_ms FROM model_splits WHERE instance_id=? AND model_id=? AND (idx > ? OR (idx = ? AND key > ?)) ORDER BY idx, key LIMIT ?",
-		c.instanceID,
-		modelID,
-		afterIndex,
-		afterIndex,
-		afterKey,
-		limit,
-	)
+func (c *catalogStore) FindModelSplits(ctx context.Context, opts *drivers.FindModelSplitsOptions) ([]drivers.ModelSplit, error) {
+	var qry strings.Builder
+	var args []any
+
+	qry.WriteString("SELECT key, data_json, idx, watermark, executed_on, error, elapsed_ms FROM model_splits WHERE instance_id=? AND model_id=?")
+	args = append(args, c.instanceID, opts.ModelID)
+
+	if opts.WhereErrored {
+		qry.WriteString(" AND error != ''")
+	}
+
+	if opts.WherePending {
+		qry.WriteString(" AND executed_on IS NULL")
+	}
+
+	if opts.AfterIndex != 0 || opts.AfterKey != "" {
+		qry.WriteString(" AND (idx > ? OR (idx = ? AND key > ?))")
+		args = append(args, opts.AfterIndex, opts.AfterIndex, opts.AfterKey)
+	}
+
+	qry.WriteString(" ORDER BY idx, key")
+
+	if opts.Limit != 0 {
+		qry.WriteString(" LIMIT ?")
+		args = append(args, opts.Limit)
+	}
+
+	rows, err := c.db.QueryContext(ctx, qry.String(), args...)
 	if err != nil {
 		return nil, err
 	}
@@ -220,38 +237,6 @@ func (c *catalogStore) FindModelSplitsByKeys(ctx context.Context, modelID string
 	return res, nil
 }
 
-func (c *catalogStore) FindModelSplitsByPending(ctx context.Context, modelID string, limit int) ([]drivers.ModelSplit, error) {
-	rows, err := c.db.QueryxContext(
-		ctx,
-		"SELECT key, data_json, idx, watermark, executed_on, error, elapsed_ms FROM model_splits WHERE instance_id=? AND model_id=? AND executed_on IS NULL ORDER BY idx LIMIT ?",
-		c.instanceID,
-		modelID,
-		limit,
-	)
-	if err != nil {
-		return nil, err
-	}
-	defer rows.Close()
-
-	var res []drivers.ModelSplit
-	for rows.Next() {
-		var elapsedMs int64
-		r := drivers.ModelSplit{}
-		err := rows.Scan(&r.Key, &r.DataJSON, &r.Index, &r.Watermark, &r.ExecutedOn, &r.Error, &elapsedMs)
-		if err != nil {
-			return nil, err
-		}
-		r.Elapsed = time.Duration(elapsedMs) * time.Millisecond
-		res = append(res, r)
-	}
-
-	if rows.Err() != nil {
-		return nil, err
-	}
-
-	return res, nil
-}
-
 func (c *catalogStore) CheckModelSplitsHaveErrors(ctx context.Context, modelID string) (bool, error) {
 	rows, err := c.db.QueryContext(
 		ctx,

runtime/reconcilers/model.go

@@ -902,7 +902,11 @@ func (r *ModelReconciler) executeAll(ctx context.Context, self *runtimev1.Resour
 	// This enables the first split to create the initial result (such as a table) that the other splits incrementally build upon.
 	if !incrementalRun {
 		// Find the first split
-		splits, err := catalog.FindModelSplitsByPending(ctx, model.State.SplitsModelId, 1)
+		splits, err := catalog.FindModelSplits(ctx, &drivers.FindModelSplitsOptions{
+			ModelID:      model.State.SplitsModelId,
+			WherePending: true,
+			Limit:        1,
+		})
 		if err != nil {
 			return "", nil, fmt.Errorf("failed to load first split: %w", err)
 		}
@@ -930,7 +934,11 @@ func (r *ModelReconciler) executeAll(ctx context.Context, self *runtimev1.Resour
 		// Get a batch of pending splits
 		// Note: We do this when no workers are running because splits are considered pending if they have not completed execution yet.
 		// This reduces concurrency when processing the last handful of splits in each batch, but with large batch sizes it's worth the simplicity for now.
-		splits, err := catalog.FindModelSplitsByPending(ctx, model.State.SplitsModelId, _modelPendingSplitsBatchSize)
+		splits, err := catalog.FindModelSplits(ctx, &drivers.FindModelSplitsOptions{
+			ModelID:      model.State.SplitsModelId,
+			WherePending: true,
+			Limit:        _modelPendingSplitsBatchSize,
+		})
 		if err != nil {
 			return "", nil, err
 		}

runtime/server/controller.go

@@ -288,7 +288,7 @@ func (s *Server) GetModelSplits(ctx context.Context, req *runtimev1.GetModelSpli
 		return &runtimev1.GetModelSplitsResponse{}, nil
 	}
 
-	afterIdx := -1
+	afterIdx := 0
 	afterKey := ""
 	if req.PageToken != "" {
 		err := unmarshalPageToken(req.PageToken, &afterIdx, &afterKey)
@@ -303,7 +303,16 @@ func (s *Server) GetModelSplits(ctx context.Context, req *runtimev1.GetModelSpli
 	}
 	defer release()
 
-	splits, err := catalog.FindModelSplits(ctx, splitsModelID, afterIdx, afterKey, validPageSize(req.PageSize))
+	opts := &drivers.FindModelSplitsOptions{
+		ModelID:      splitsModelID,
+		WherePending: req.Pending,
+		WhereErrored: req.Errored,
+		AfterIndex:   afterIdx,
+		AfterKey:     afterKey,
+		Limit:        validPageSize(req.PageSize),
+	}
+
+	splits, err := catalog.FindModelSplits(ctx, opts)
 	if err != nil {
 		return nil, status.Error(codes.InvalidArgument, err.Error())
 	}

web-common/src/proto/gen/rill/runtime/v1/api_pb.ts

@@ -3092,6 +3092,16 @@ export class GetModelSplitsRequest extends Message<GetModelSplitsRequest> {
    */
   model = "";
 
+  /**
+   * @generated from field: bool pending = 5;
+   */
+  pending = false;
+
+  /**
+   * @generated from field: bool errored = 6;
+   */
+  errored = false;
+
   /**
    * @generated from field: uint32 page_size = 3;
    */
@@ -3112,6 +3122,8 @@ export class GetModelSplitsRequest extends Message<GetModelSplitsRequest> {
   static readonly fields: FieldList = proto3.util.newFieldList(() => [
     { no: 1, name: "instance_id", kind: "scalar", T: 9 /* ScalarType.STRING */ },
     { no: 2, name: "model", kind: "scalar", T: 9 /* ScalarType.STRING */ },
+    { no: 5, name: "pending", kind: "scalar", T: 8 /* ScalarType.BOOL */ },
+    { no: 6, name: "errored", kind: "scalar", T: 8 /* ScalarType.BOOL */ },
     { no: 3, name: "page_size", kind: "scalar", T: 13 /* ScalarType.UINT32 */ },
     { no: 4, name: "page_token", kind: "scalar", T: 9 /* ScalarType.STRING */ },
   ]);

web-common/src/runtime-client/gen/index.schemas.ts

@@ -343,6 +343,8 @@ export type QueryServiceColumnCardinalityParams = {
 };
 
 export type RuntimeServiceGetModelSplitsParams = {
+  pending?: boolean;
+  errored?: boolean;
   pageSize?: number;
   pageToken?: string;
 };