Implement CHASM List/Count Executions

Author: awln-temporal

chasm/engine.go

@@ -4,6 +4,11 @@ package chasm
 
 import (
 	"context"
+	"reflect"
+	"time"
+
+	commonpb "go.temporal.io/api/common/v1"
+	"google.golang.org/protobuf/proto"
 )
 
 // NoValue is a sentinel type representing no value.
@@ -46,6 +51,59 @@ type Engine interface {
 		func(MutableContext, Component, any) error,
 		...TransitionOption,
 	) ([]byte, error)
+
+	ListExecutions(
+		context.Context,
+		reflect.Type,
+		*ListExecutionsRequest,
+	) (*ListExecutionsResponse[*commonpb.Payload], error)
+
+	CountExecutions(
+		context.Context,
+		reflect.Type,
+		*CountExecutionsRequest,
+	) (*CountExecutionsResponse, error)
+}
+
+type ListExecutionsResponse[M proto.Message] struct {
+	Executions    []*ExecutionInfo[M]
+	NextPageToken []byte
+}
+type ExecutionInfo[M proto.Message] struct {
+	BusinessID             string
+	RunID                  string
+	StartTime              time.Time
+	CloseTime              time.Time
+	HistoryLength          int64
+	HistorySizeBytes       int64
+	StateTransitionCount   int64
+	ChasmSearchAttributes  SearchAttributesMap
+	CustomSearchAttributes map[string]*commonpb.Payload
+	Memo                   *commonpb.Memo
+	ChasmMemo              M
+}
+
+type ListExecutionsRequest struct {
+	NamespaceID   string
+	NamespaceName string
+	Query         string
+	PageSize      int
+	NextPageToken []byte
+}
+
+type ListExecutionsOptions struct {
+	PageSize      int
+	NextPageToken []byte
+}
+
+type CountExecutionsRequest struct {
+	NamespaceID   string
+	NamespaceName string
+	Query         string
+}
+
+type CountExecutionsResponse struct {
+	Count int64
 }
 
 type BusinessIDReusePolicy int
@@ -304,3 +362,48 @@ func engineFromContext(
 	}
 	return e
 }
+
+func ListExecutions[C Component, M proto.Message](
+	ctx context.Context,
+	request *ListExecutionsRequest,
+) (*ListExecutionsResponse[M], error) {
+	archetypeType := reflect.TypeFor[C]()
+	response, err := engineFromContext(ctx).ListExecutions(ctx, archetypeType, request)
+	if err != nil {
+		return nil, err
+	}
+
+	// Convert response, unmarshaling ChasmMemo to type M
+	executions := make([]*ExecutionInfo[M], len(response.Executions))
+	for i, execution := range response.Executions {
+		chasmMemo := reflect.New(reflect.TypeFor[M]().Elem()).Interface().(M)
+		if err := proto.Unmarshal(execution.ChasmMemo.Data, chasmMemo); err == nil {
+			executions[i] = &ExecutionInfo[M]{
+				BusinessID:             execution.BusinessID,
+				RunID:                  execution.RunID,
+				StartTime:              execution.StartTime,
+				CloseTime:              execution.CloseTime,
+				HistoryLength:          execution.HistoryLength,
+				HistorySizeBytes:       execution.HistorySizeBytes,
+				StateTransitionCount:   execution.StateTransitionCount,
+				ChasmSearchAttributes:  execution.ChasmSearchAttributes,
+				CustomSearchAttributes: execution.CustomSearchAttributes,
+				Memo:                   execution.Memo,
+				ChasmMemo:              chasmMemo,
+			}
+		}
+	}
+
+	return &ListExecutionsResponse[M]{
+		Executions:    executions,
+		NextPageToken: response.NextPageToken,
+	}, nil
+}
+
+func CountExecutions[C Component](
+	ctx context.Context,
+	request *CountExecutionsRequest,
+) (*CountExecutionsResponse, error) {
+	archetypeType := reflect.TypeFor[C]()
+	return engineFromContext(ctx).CountExecutions(ctx, archetypeType, request)
+}

chasm/engine_mock.go

@@ -21,10 +21,6 @@ func (r *Registry) ComponentOf(componentGoType reflect.Type) (*RegistrableCompon
 	return r.componentOf(componentGoType)
 }
 
-func (r *Registry) ComponentByID(id uint32) (*RegistrableComponent, bool) {
-	return r.componentByID(id)
-}
-
 func (r *Registry) TaskFor(taskInstance any) (*RegistrableTask, bool) {
 	return r.taskFor(taskInstance)
 }

chasm/export_test.go

@@ -6,6 +6,7 @@ import (
 	"go.temporal.io/server/chasm"
 	"go.temporal.io/server/chasm/lib/tests/gen/testspb/v1"
 	"go.temporal.io/server/common"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
@@ -154,11 +155,6 @@ func (s *PayloadStore) SearchAttributes(
 }
 
 // Memo implements chasm.VisibilityMemoProvider interface
-func (s *PayloadStore) Memo(
-	_ chasm.Context,
-) map[string]chasm.VisibilityValue {
-	return map[string]chasm.VisibilityValue{
-		TotalCountMemoFieldName: chasm.VisibilityValueInt64(s.State.TotalCount),
-		TotalSizeMemoFieldName:  chasm.VisibilityValueInt64(s.State.TotalSize),
-	}
+func (s *PayloadStore) Memo(_ chasm.Context) proto.Message {
+	return s.State
 }

chasm/lib/tests/payload.go

@@ -51,7 +51,7 @@ func (s *componentRefSuite) TestArchetypeID() {
 	archetypeID, err := ref.ArchetypeID(s.registry)
 	s.NoError(err)
 
-	rc, ok := s.registry.ComponentOf(reflect.TypeFor[*TestComponent]())
+	rc, ok := s.registry.componentOf(reflect.TypeFor[*TestComponent]())
 	s.True(ok)
 
 	s.Equal(rc.componentID, archetypeID)
@@ -68,7 +68,7 @@ func (s *componentRefSuite) TestShardingKey() {
 	shardingKey, err := ref.ShardingKey(s.registry)
 	s.NoError(err)
 
-	rc, ok := s.registry.ComponentOf(reflect.TypeFor[*TestComponent]())
+	rc, ok := s.registry.componentOf(reflect.TypeFor[*TestComponent]())
 	s.True(ok)
 
 	s.Equal(rc.shardingFn(entityKey), shardingKey)

chasm/ref_test.go

@@ -115,6 +115,11 @@ func (rc *RegistrableComponent) registerToLibrary(
 	return fqn, rc.componentID, nil
 }
 
+// SearchAttributesMapper returns the search attributes mapper for this component.
+func (rc *RegistrableComponent) SearchAttributesMapper() *VisibilitySearchAttributesMapper {
+	return rc.searchAttributesMapper
+}
+
 // fqType returns the fully qualified name of the component, which is a combination of
 // the library name and the component type. This is used to uniquely identify
 // the component in the registry.

chasm/registrable_component.go

@@ -81,6 +81,15 @@ func (r *Registry) ComponentFqnByID(id uint32) (string, bool) {
 	return fqn, ok
 }
 
+// ComponentByID returns the registrable component for a given archetype ID.
+func (r *Registry) ComponentByID(id uint32) (*RegistrableComponent, bool) {
+	fqn, ok := r.componentFqnByID[id]
+	if !ok {
+		return nil, false
+	}
+	return r.component(fqn)
+}
+
 // ComponentIDFor converts registered component instance to component type ID.
 // This method should only be used by CHASM framework internal code,
 // NOT CHASM library developers.
@@ -136,6 +145,15 @@ func (r *Registry) componentOf(componentGoType reflect.Type) (*RegistrableCompon
 	return rc, ok
 }
 
+// ArchetypeIDOf returns the ArchetypeID for the given component Go type.
+func (r *Registry) ArchetypeIDOf(componentGoType reflect.Type) (ArchetypeID, bool) {
+	rc, ok := r.componentByGoType[componentGoType]
+	if !ok {
+		return UnspecifiedArchetypeID, false
+	}
+	return rc.componentID, true
+}
+
 func (r *Registry) taskOf(taskGoType reflect.Type) (*RegistrableTask, bool) {
 	rt, ok := r.taskByGoType[taskGoType]
 	return rt, ok