Address remaining commments

Author: awln-temporal

chasm/search_attribute.go

@@ -3,6 +3,7 @@ package chasm
 import (
 	"errors"
 	"fmt"
+	"reflect"
 	"time"
 
 	enumspb "go.temporal.io/api/enums/v1"
@@ -67,13 +68,19 @@ var (
 )
 
 var (
-	_ SearchAttribute = (*searchAttributeDefinition)(nil)
 	_ SearchAttribute = (*SearchAttributeBool)(nil)
 	_ SearchAttribute = (*SearchAttributeDateTime)(nil)
 	_ SearchAttribute = (*SearchAttributeInt)(nil)
 	_ SearchAttribute = (*SearchAttributeDouble)(nil)
 	_ SearchAttribute = (*SearchAttributeKeyword)(nil)
 	_ SearchAttribute = (*SearchAttributeKeywordList)(nil)
+
+	_ typedSearchAttribute[bool]      = (*SearchAttributeBool)(nil)
+	_ typedSearchAttribute[time.Time] = (*SearchAttributeDateTime)(nil)
+	_ typedSearchAttribute[int64]     = (*SearchAttributeInt)(nil)
+	_ typedSearchAttribute[float64]   = (*SearchAttributeDouble)(nil)
+	_ typedSearchAttribute[string]    = (*SearchAttributeKeyword)(nil)
+	_ typedSearchAttribute[[]string]  = (*SearchAttributeKeywordList)(nil)
 )
 
 type (
@@ -82,6 +89,12 @@ type (
 		definition() searchAttributeDefinition
 	}
 
+	typedSearchAttribute[T any] interface {
+		SearchAttribute
+		typeMarker(T)
+		markedVisibilityType() reflect.Type
+	}
+
 	searchAttributeDefinition struct {
 		alias     string
 		field     string
@@ -210,6 +223,12 @@ func (s SearchAttributeBool) Value(value bool) SearchAttributeKeyValue {
 	}
 }
 
+func (s SearchAttributeBool) typeMarker(_ bool) {}
+
+func (s SearchAttributeBool) markedVisibilityType() reflect.Type {
+	return reflect.TypeFor[VisibilityValueBool]()
+}
+
 // SearchAttributeDateTime is a search attribute for a datetime value.
 type SearchAttributeDateTime struct {
 	searchAttributeDefinition
@@ -245,6 +264,12 @@ func (s SearchAttributeDateTime) Value(value time.Time) SearchAttributeKeyValue
 	}
 }
 
+func (s SearchAttributeDateTime) typeMarker(_ time.Time) {}
+
+func (s SearchAttributeDateTime) markedVisibilityType() reflect.Type {
+	return reflect.TypeFor[VisibilityValueTime]()
+}
+
 // SearchAttributeInt is a search attribute for an integer value.
 type SearchAttributeInt struct {
 	searchAttributeDefinition
@@ -280,6 +305,12 @@ func (s SearchAttributeInt) Value(value int64) SearchAttributeKeyValue {
 	}
 }
 
+func (s SearchAttributeInt) typeMarker(_ int64) {}
+
+func (s SearchAttributeInt) markedVisibilityType() reflect.Type {
+	return reflect.TypeFor[VisibilityValueInt64]()
+}
+
 // SearchAttributeDouble is a search attribute for a double value.
 type SearchAttributeDouble struct {
 	searchAttributeDefinition
@@ -315,6 +346,12 @@ func (s SearchAttributeDouble) Value(value float64) SearchAttributeKeyValue {
 	}
 }
 
+func (s SearchAttributeDouble) typeMarker(_ float64) {}
+
+func (s SearchAttributeDouble) markedVisibilityType() reflect.Type {
+	return reflect.TypeFor[VisibilityValueFloat64]()
+}
+
 // SearchAttributeKeyword is a search attribute for a keyword value.
 type SearchAttributeKeyword struct {
 	searchAttributeDefinition
@@ -350,6 +387,12 @@ func (s SearchAttributeKeyword) Value(value string) SearchAttributeKeyValue {
 	}
 }
 
+func (s SearchAttributeKeyword) typeMarker(_ string) {}
+
+func (s SearchAttributeKeyword) markedVisibilityType() reflect.Type {
+	return reflect.TypeFor[VisibilityValueString]()
+}
+
 // SearchAttributeKeywordList is a search attribute for a keyword list value.
 type SearchAttributeKeywordList struct {
 	searchAttributeDefinition
@@ -385,6 +428,12 @@ func (s SearchAttributeKeywordList) Value(value []string) SearchAttributeKeyValu
 	}
 }
 
+func (s SearchAttributeKeywordList) typeMarker(_ []string) {}
+
+func (s SearchAttributeKeywordList) markedVisibilityType() reflect.Type {
+	return reflect.TypeFor[VisibilityValueStringSlice]()
+}
+
 // SearchAttributesMap wraps search attribute values with type-safe access.
 type SearchAttributesMap struct {
 	values map[string]VisibilityValue
@@ -395,94 +444,39 @@ func NewSearchAttributesMap(values map[string]VisibilityValue) SearchAttributesM
 	return SearchAttributesMap{values: values}
 }
 
-// GetBool returns the boolean value for a given SearchAttributeBool. If not found or map is nil, found is false.
-func (m SearchAttributesMap) GetBool(sa SearchAttributeBool) (value bool, found bool) {
+// Get returns the value for a given SearchAttribute with compile-time type safety.
+// The return type T is inferred from the SearchAttribute's type parameter.
+// For example, SearchAttriteBool will return a bool value.
+// If the value is not found, the zero value for the type T is returned and the second return value is false.
+func Get[T any](m SearchAttributesMap, sa typedSearchAttribute[T]) (val T, ok bool) {
+	var zero T
 	if m.values == nil {
-		return false, false
+		return zero, false
 	}
 
 	alias := sa.definition().alias
-	boolVal, ok := m.values[alias].(VisibilityValueBool)
-	if !ok {
-		return false, false
+	visibilityValue, exists := m.values[alias]
+	if !exists {
+		return zero, false
 	}
 
-	return bool(boolVal), true
-}
-
-// GetInt returns the int value for a given SearchAttributeInt. If not found or map is nil, second parameter is false.
-func (m SearchAttributesMap) GetInt(sa SearchAttributeInt) (int, bool) {
-	if m.values == nil {
-		return 0, false
+	if reflect.TypeOf(visibilityValue) != sa.markedVisibilityType() {
+		return zero, false
 	}
 
-	alias := sa.definition().alias
-	intValue, ok := m.values[alias].(VisibilityValueInt)
-	if !ok {
-		return 0, false
-	}
-
-	return int(intValue), true
-}
+	reflectVal := reflect.ValueOf(val)
+	targetType := reflect.TypeFor[T]()
 
-// GetDouble returns the double value for a given SearchAttributeDouble. If not found or map is nil, second parameter is false.
-func (m SearchAttributesMap) GetDouble(sa SearchAttributeDouble) (float64, bool) {
-	if m.values == nil {
-		return 0, false
-	}
-
-	alias := sa.definition().alias
-	doubleValue, ok := m.values[alias].(VisibilityValueFloat64)
-	if !ok {
-		return 0, false
+	if !reflectVal.Type().ConvertibleTo(targetType) {
+		return zero, false
 	}
 
-	return float64(doubleValue), true
-}
-
-// GetKeyword returns the string value for a given SearchAttributeKeyword. If not found or map is nil, second parameter is false.
-func (m SearchAttributesMap) GetKeyword(sa SearchAttributeKeyword) (string, bool) {
-	if m.values == nil {
-		return "", false
-	}
-
-	alias := sa.definition().alias
-	stringValue, ok := m.values[alias].(VisibilityValueString)
-	if !ok {
-		return "", false
-	}
-
-	return string(stringValue), true
-}
-
-// GetDateTime returns the time value for a given SearchAttributeDateTime. If not found or map is nil, second parameter is false.
-func (m SearchAttributesMap) GetDateTime(sa SearchAttributeDateTime) (time.Time, bool) {
-	if m.values == nil {
-		return time.Time{}, false
-	}
-
-	alias := sa.definition().alias
-	timeValue, ok := m.values[alias].(VisibilityValueTime)
-	if !ok {
-		return time.Time{}, false
-	}
-
-	return time.Time(timeValue), true
-}
-
-// GetKeywordList returns the string list value for a given SearchAttributeKeywordList. If not found or map is nil, second parameter is false.
-func (m SearchAttributesMap) GetKeywordList(sa SearchAttributeKeywordList) ([]string, bool) {
-	if m.values == nil {
-		return nil, false
-	}
-
-	alias := sa.definition().alias
-	keywordListValue, ok := m.values[alias].(VisibilityValueStringSlice)
+	result, ok := reflectVal.Convert(targetType).Interface().(T)
 	if !ok {
-		return nil, false
+		return zero, false
 	}
 
-	return []string(keywordListValue), true
+	return result, true
 }
 
 // convertToVisibilityValue converts a value to VisibilityValue based on its runtime type.

chasm/search_attribute_test.go

@@ -0,0 +1,81 @@
+package chasm
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestSearchAttributesMap_Get(t *testing.T) {
+	// Define test search attributes
+	boolAttr := NewSearchAttributeBool("completed", SearchAttributeFieldBool01)
+	intAttr := NewSearchAttributeInt("count", SearchAttributeFieldInt01)
+	doubleAttr := NewSearchAttributeDouble("score", SearchAttributeFieldDouble01)
+	keywordAttr := NewSearchAttributeKeyword("status", SearchAttributeFieldKeyword01)
+	datetimeAttr := NewSearchAttributeDateTime("timestamp", SearchAttributeFieldDateTime01)
+	keywordListAttr := NewSearchAttributeKeywordList("tags", SearchAttributeFieldKeywordList01)
+
+	now := time.Now()
+
+	// Create map with test values
+	values := map[string]VisibilityValue{
+		"completed": VisibilityValueBool(true),
+		"count":     VisibilityValueInt64(42),
+		"score":     VisibilityValueFloat64(3.14),
+		"status":    VisibilityValueString("active"),
+		"timestamp": VisibilityValueTime(now),
+		"tags":      VisibilityValueStringSlice([]string{"tag1", "tag2"}),
+	}
+	m := NewSearchAttributesMap(values)
+
+	t.Run("GetBool", func(t *testing.T) {
+		val, ok := Get(m, boolAttr)
+		assert.True(t, ok)
+		assert.True(t, val)
+	})
+
+	t.Run("GetInt64", func(t *testing.T) {
+		val, ok := Get(m, intAttr)
+		assert.True(t, ok)
+		assert.Equal(t, int64(42), val)
+	})
+
+	t.Run("GetFloat64", func(t *testing.T) {
+		val, ok := Get(m, doubleAttr)
+		assert.True(t, ok)
+		assert.InDelta(t, 3.14, val, 0.0001)
+	})
+
+	t.Run("GetString", func(t *testing.T) {
+		val, ok := Get(m, keywordAttr)
+		assert.True(t, ok)
+		assert.Equal(t, "active", val)
+	})
+
+	t.Run("GetTime", func(t *testing.T) {
+		val, ok := Get(m, datetimeAttr)
+		assert.True(t, ok)
+		assert.True(t, now.Equal(val))
+	})
+
+	t.Run("GetStringSlice", func(t *testing.T) {
+		val, ok := Get(m, keywordListAttr)
+		assert.True(t, ok)
+		assert.Equal(t, []string{"tag1", "tag2"}, val)
+	})
+
+	t.Run("NotFound", func(t *testing.T) {
+		missingAttr := NewSearchAttributeBool("missing", SearchAttributeFieldBool02)
+		val, ok := Get(m, missingAttr)
+		assert.False(t, ok)
+		assert.False(t, val)
+	})
+
+	t.Run("NilMap", func(t *testing.T) {
+		emptyMap := NewSearchAttributesMap(nil)
+		val, ok := Get(emptyMap, boolAttr)
+		assert.False(t, ok)
+		assert.False(t, val)
+	})
+}

common/persistence/visibility/manager/visibility_manager.go

@@ -36,7 +36,7 @@ type (
 		ListWorkflowExecutions(ctx context.Context, request *ListWorkflowExecutionsRequestV2) (*ListWorkflowExecutionsResponse, error)
 		ListChasmExecutions(ctx context.Context, request *ListChasmExecutionsRequest) (*chasm.ListExecutionsResponse[*commonpb.Payload], error)
 		CountWorkflowExecutions(ctx context.Context, request *CountWorkflowExecutionsRequest) (*CountWorkflowExecutionsResponse, error)
-		CountChasmExecutions(ctx context.Context, request *CountChasmExecutionsRequest) (*CountChasmExecutionsResponse, error)
+		CountChasmExecutions(ctx context.Context, request *CountChasmExecutionsRequest) (*chasm.CountExecutionsResponse, error)
 		GetWorkflowExecution(ctx context.Context, request *GetWorkflowExecutionRequest) (*GetWorkflowExecutionResponse, error)
 
 		// Admin APIs

common/persistence/visibility/manager/visibility_manager_mock.go

@@ -190,7 +190,7 @@ func (v *VisibilityManagerDual) CountWorkflowExecutions(
 func (v *VisibilityManagerDual) CountChasmExecutions(
 	ctx context.Context,
 	request *manager.CountChasmExecutionsRequest,
-) (*manager.CountChasmExecutionsResponse, error) {
+) (*chasm.CountExecutionsResponse, error) {
 	return dualReadWrapper(
 		ctx,
 		v,

common/persistence/visibility/visibility_manager_dual.go

@@ -211,13 +211,13 @@ func (p *visibilityManagerImpl) ListChasmExecutions(
 func (p *visibilityManagerImpl) CountChasmExecutions(
 	ctx context.Context,
 	request *manager.CountChasmExecutionsRequest,
-) (*manager.CountChasmExecutionsResponse, error) {
+) (*chasm.CountExecutionsResponse, error) {
 	response, err := p.store.CountChasmExecutions(ctx, request)
 	if err != nil {
 		return nil, err
 	}
 
-	return response, nil
+	return &chasm.CountExecutionsResponse{Count: response.Count}, nil
 }
 
 func (p *visibilityManagerImpl) CountWorkflowExecutions(

common/persistence/visibility/visibility_manager_impl.go

@@ -141,7 +141,7 @@ func (m *visibilityManagerRateLimited) CountWorkflowExecutions(
 func (m *visibilityManagerRateLimited) CountChasmExecutions(
 	ctx context.Context,
 	request *manager.CountChasmExecutionsRequest,
-) (*manager.CountChasmExecutionsResponse, error) {
+) (*chasm.CountExecutionsResponse, error) {
 	if ok := allow(ctx, "CountChasmExecutions", m.readRateLimiter); !ok {
 		return nil, persistence.ErrPersistenceSystemLimitExceeded
 	}

common/persistence/visibility/visibility_manager_rate_limited.go

@@ -162,7 +162,7 @@ func (m *visibilityManagerMetrics) CountWorkflowExecutions(
 func (m *visibilityManagerMetrics) CountChasmExecutions(
 	ctx context.Context,
 	request *manager.CountChasmExecutionsRequest,
-) (*manager.CountChasmExecutionsResponse, error) {
+) (*chasm.CountExecutionsResponse, error) {
 	handler, startTime := m.tagScope(metrics.VisibilityPersistenceCountChasmExecutionsScope)
 	response, err := m.delegate.CountChasmExecutions(ctx, request)
 	metrics.VisibilityPersistenceLatency.With(handler).Record(time.Since(startTime))

common/persistence/visibility/visiblity_manager_metrics.go

@@ -670,15 +670,10 @@ func (e *ChasmEngine) CountExecutions(
 		return nil, serviceerror.NewInternal("unknown chasm component type: " + archetypeType.String())
 	}
 
-	resp, err := e.visibilityMgr.CountChasmExecutions(ctx, &manager.CountChasmExecutionsRequest{
+	return e.visibilityMgr.CountChasmExecutions(ctx, &manager.CountChasmExecutionsRequest{
 		ArchetypeID: archetypeID,
 		NamespaceID: namespace.ID(request.NamespaceID),
 		Namespace:   namespace.Name(request.NamespaceName),
 		Query:       request.Query,
 	})
-	if err != nil {
-		return nil, err
-	}
-
-	return &chasm.CountExecutionsResponse{Count: resp.Count}, nil
 }