fix: mark idempotent CQL queries for driver-level retry on write timeouts

common/persistence/cassandra/cluster_metadata_store.go

@@ -64,7 +64,7 @@ func (m *ClusterMetadataStore) ListClusterMetadata(
 	ctx context.Context,
 	request *p.InternalListClusterMetadataRequest,
 ) (*p.InternalListClusterMetadataResponse, error) {
-	query := m.session.Query(templateListClusterMetadata, constMetadataPartition).WithContext(ctx)
+	query := m.session.Query(templateListClusterMetadata, constMetadataPartition).WithContext(ctx).Idempotent(true)
 	iter := query.PageSize(request.PageSize).PageState(request.NextPageToken).Iter()
 
 	response := &p.InternalListClusterMetadataResponse{}
@@ -106,7 +106,7 @@ func (m *ClusterMetadataStore) GetClusterMetadata(
 	var encoding string
 	var version int64
 
-	query := m.session.Query(templateGetClusterMetadata, constMetadataPartition, request.ClusterName).WithContext(ctx)
+	query := m.session.Query(templateGetClusterMetadata, constMetadataPartition, request.ClusterName).WithContext(ctx).Idempotent(true)
 	err := query.Scan(&clusterMetadata, &encoding, &version)
 	if err != nil {
 		return nil, gocql.ConvertError("GetClusterMetadata", err)
@@ -159,7 +159,7 @@ func (m *ClusterMetadataStore) DeleteClusterMetadata(
 	ctx context.Context,
 	request *p.InternalDeleteClusterMetadataRequest,
 ) error {
-	query := m.session.Query(templateDeleteClusterMetadata, constMetadataPartition, request.ClusterName).WithContext(ctx)
+	query := m.session.Query(templateDeleteClusterMetadata, constMetadataPartition, request.ClusterName).WithContext(ctx).Idempotent(true)
 	if err := query.Exec(); err != nil {
 		return gocql.ConvertError("DeleteClusterMetadata", err)
 	}
@@ -202,7 +202,7 @@ func (m *ClusterMetadataStore) GetClusterMembers(
 	}
 
 	queryString.WriteString(templateAllowFiltering)
-	query := m.session.Query(queryString.String(), operands...).WithContext(ctx)
+	query := m.session.Query(queryString.String(), operands...).WithContext(ctx).Idempotent(true)
 
 	iter := query.PageSize(request.PageSize).PageState(request.NextPageToken).Iter()
 
@@ -257,7 +257,7 @@ func (m *ClusterMetadataStore) UpsertClusterMembership(
 		request.SessionStart,
 		time.Now().UTC(),
 		int64(request.RecordExpiry.Seconds()),
-	).WithContext(ctx)
+	).WithContext(ctx).Idempotent(true)
 	err := query.Exec()
 
 	if err != nil {

common/persistence/cassandra/history_store.go

@@ -77,7 +77,7 @@ func (h *HistoryStore) AppendHistoryNodes(
 			node.TransactionID,
 			node.Events.Data,
 			node.Events.EncodingType.String(),
-		).WithContext(ctx)
+		).WithContext(ctx).Idempotent(true)
 		if err := query.Exec(); err != nil {
 			return convertTimeoutError(gocql.ConvertError("AppendHistoryNodes", err))
 		}
@@ -129,7 +129,7 @@ func (h *HistoryStore) DeleteHistoryNodes(
 		branchID,
 		nodeID,
 		txnID,
-	).WithContext(ctx)
+	).WithContext(ctx).Idempotent(true)
 	if err := query.Exec(); err != nil {
 		return gocql.ConvertError("DeleteHistoryNodes", err)
 	}
@@ -166,7 +166,7 @@ func (h *HistoryStore) ReadHistoryBranch(
 		queryString = v2templateReadHistoryNode
 	}
 
-	query := h.Session.Query(queryString, treeID, branchID, request.MinNodeID, request.MaxNodeID).WithContext(ctx)
+	query := h.Session.Query(queryString, treeID, branchID, request.MinNodeID, request.MaxNodeID).WithContext(ctx).Idempotent(true)
 
 	iter := query.PageSize(request.PageSize).PageState(request.NextPageToken).Iter()
 	var pagingToken []byte
@@ -255,7 +255,7 @@ func (h *HistoryStore) ForkHistoryBranch(
 	if err != nil {
 		return serviceerror.NewInternalf("ForkHistoryBranch - Gocql NewBranchID UUID cast failed. Error: %v", err)
 	}
-	query := h.Session.Query(v2templateInsertTree, cqlTreeID, cqlNewBranchID, datablob.Data, datablob.EncodingType.String()).WithContext(ctx)
+	query := h.Session.Query(v2templateInsertTree, cqlTreeID, cqlNewBranchID, datablob.Data, datablob.EncodingType.String()).WithContext(ctx).Idempotent(true)
 	err = query.Exec()
 	if err != nil {
 		return gocql.ConvertError("ForkHistoryBranch", err)
@@ -303,7 +303,7 @@ func (h *HistoryStore) GetAllHistoryTreeBranches(
 	request *p.GetAllHistoryTreeBranchesRequest,
 ) (*p.InternalGetAllHistoryTreeBranchesResponse, error) {
 
-	query := h.Session.Query(v2templateScanAllTreeBranches).WithContext(ctx)
+	query := h.Session.Query(v2templateScanAllTreeBranches).WithContext(ctx).Idempotent(true)
 
 	iter := query.PageSize(request.PageSize).PageState(request.NextPageToken).Iter()
 
@@ -360,7 +360,7 @@ func (h *HistoryStore) GetHistoryTreeContainingBranch(
 	if err != nil {
 		return nil, serviceerror.NewInternalf("ReadHistoryBranch. Gocql TreeId UUID cast failed. Error: %v", err)
 	}
-	query := h.Session.Query(v2templateReadAllBranches, treeID).WithContext(ctx)
+	query := h.Session.Query(v2templateReadAllBranches, treeID).WithContext(ctx).Idempotent(true)
 
 	pageSize := 100
 	var pagingToken []byte

common/persistence/cassandra/matching_task_store_queue.go

@@ -163,7 +163,7 @@ func (d *taskQueueStore) GetTaskQueue(
 		request.TaskType,
 		rowTypeTaskQueue,
 		taskQueueTaskID,
-	).WithContext(ctx)
+	).WithContext(ctx).Idempotent(true)
 
 	var rangeID int64
 	var tlBytes []byte

common/persistence/cassandra/matching_task_store_user_data.go

@@ -48,7 +48,7 @@ func (d *userDataStore) GetTaskQueueUserData(
 	query := d.Session.Query(templateGetTaskQueueUserDataQuery,
 		request.NamespaceID,
 		request.TaskQueue,
-	).WithContext(ctx)
+	).WithContext(ctx).Idempotent(true)
 	var version int64
 	var userDataBytes []byte
 	var encoding string
@@ -134,7 +134,7 @@ func (d *userDataStore) UpdateTaskQueueUserData(
 }
 
 func (d *userDataStore) ListTaskQueueUserDataEntries(ctx context.Context, request *p.ListTaskQueueUserDataEntriesRequest) (*p.InternalListTaskQueueUserDataEntriesResponse, error) {
-	query := d.Session.Query(templateListTaskQueueUserDataQuery, request.NamespaceID).WithContext(ctx)
+	query := d.Session.Query(templateListTaskQueueUserDataQuery, request.NamespaceID).WithContext(ctx).Idempotent(true)
 	iter := query.PageSize(request.PageSize).PageState(request.NextPageToken).Iter()
 
 	response := &p.InternalListTaskQueueUserDataEntriesResponse{}
@@ -172,7 +172,7 @@ func (d *userDataStore) ListTaskQueueUserDataEntries(ctx context.Context, reques
 }
 
 func (d *userDataStore) GetTaskQueuesByBuildId(ctx context.Context, request *p.GetTaskQueuesByBuildIdRequest) ([]string, error) {
-	query := d.Session.Query(templateListTaskQueueNamesByBuildIdQuery, request.NamespaceID, request.BuildID).WithContext(ctx)
+	query := d.Session.Query(templateListTaskQueueNamesByBuildIdQuery, request.NamespaceID, request.BuildID).WithContext(ctx).Idempotent(true)
 	iter := query.PageSize(listTaskQueueNamesByBuildIdPageSize).Iter()
 
 	var taskQueues []string
@@ -207,7 +207,7 @@ func (d *userDataStore) GetTaskQueuesByBuildId(ctx context.Context, request *p.G
 
 func (d *userDataStore) CountTaskQueuesByBuildId(ctx context.Context, request *p.CountTaskQueuesByBuildIdRequest) (int, error) {
 	var count int
-	query := d.Session.Query(templateCountTaskQueueByBuildIdQuery, request.NamespaceID, request.BuildID).WithContext(ctx)
+	query := d.Session.Query(templateCountTaskQueueByBuildIdQuery, request.NamespaceID, request.BuildID).WithContext(ctx).Idempotent(true)
 	err := query.Scan(&count)
 	return count, err
 }

common/persistence/cassandra/matching_task_store_v1.go

@@ -150,7 +150,7 @@ func (d *matchingTaskStoreV1) GetTasks(
 		rowTypeTaskInSubqueue(request.Subqueue),
 		request.InclusiveMinTaskID,
 		request.ExclusiveMaxTaskID,
-	).WithContext(ctx)
+	).WithContext(ctx).Idempotent(true)
 	iter := query.PageSize(request.PageSize).PageState(request.NextPageToken).Iter()
 
 	response := &p.InternalGetTasksResponse{}
@@ -212,7 +212,7 @@ func (d *matchingTaskStoreV1) CompleteTasksLessThan(
 		request.TaskType,
 		rowTypeTaskInSubqueue(request.Subqueue),
 		request.ExclusiveMaxTaskID,
-	).WithContext(ctx)
+	).WithContext(ctx).Idempotent(true)
 	err := query.Exec()
 	if err != nil {
 		return 0, gocql.ConvertError("CompleteTasksLessThan", err)

common/persistence/cassandra/matching_task_store_v2.go

@@ -169,7 +169,7 @@ func (d *matchingTaskStoreV2) GetTasks(
 			int64(math.MaxInt64),
 		)
 	}
-	iter := query.WithContext(ctx).PageSize(request.PageSize).PageState(request.NextPageToken).Iter()
+	iter := query.WithContext(ctx).Idempotent(true).PageSize(request.PageSize).PageState(request.NextPageToken).Iter()
 
 	response := &p.InternalGetTasksResponse{}
 	task := make(map[string]any)
@@ -236,7 +236,7 @@ func (d *matchingTaskStoreV2) CompleteTasksLessThan(
 		rowType,
 		request.ExclusiveMaxPass,
 		request.ExclusiveMaxTaskID,
-	).WithContext(ctx)
+	).WithContext(ctx).Idempotent(true)
 	err := query.Exec()
 	if err != nil {
 		return 0, gocql.ConvertError("CompleteTasksLessThan", err)

common/persistence/cassandra/metadata_store.go

@@ -154,7 +154,7 @@ func (m *MetadataStore) CreateNamespaceInV2Table(
 	defer func() { _ = iter.Close() }()
 	deleteOrphanNamespace := func() {
 		// Delete namespace from `namespaces_by_id`
-		if errDelete := m.session.Query(templateDeleteNamespaceQuery, request.ID).WithContext(ctx).Exec(); errDelete != nil {
+		if errDelete := m.session.Query(templateDeleteNamespaceQuery, request.ID).WithContext(ctx).Idempotent(true).Exec(); errDelete != nil {
 			m.logger.Warn("Unable to delete orphan namespace record. Error", tag.Error(errDelete))
 		}
 	}
@@ -251,7 +251,7 @@ func (m *MetadataStore) RenameNamespace(
 	if updateErr := m.session.Query(templateUpdateNamespaceByIdQuery,
 		request.Name,
 		request.Id,
-	).WithContext(ctx).Exec(); updateErr != nil {
+	).WithContext(ctx).Idempotent(true).Exec(); updateErr != nil {
 		return gocql.ConvertError("RenameNamespace", updateErr)
 	}
 
@@ -316,14 +316,14 @@ func (m *MetadataStore) GetNamespace(
 
 	namespace := request.Name
 	if len(request.ID) > 0 {
-		query = m.session.Query(templateGetNamespaceQuery, request.ID).WithContext(ctx)
+		query = m.session.Query(templateGetNamespaceQuery, request.ID).WithContext(ctx).Idempotent(true)
 		err = query.Scan(&namespace)
 		if err != nil {
 			return nil, handleError(request.Name, request.ID, err)
 		}
 	}
 
-	query = m.session.Query(templateGetNamespaceByNameQueryV2, constNamespacePartition, namespace).WithContext(ctx)
+	query = m.session.Query(templateGetNamespaceByNameQueryV2, constNamespacePartition, namespace).WithContext(ctx).Idempotent(true)
 	err = query.Scan(
 		nil,
 		nil,
@@ -352,7 +352,7 @@ func (m *MetadataStore) ListNamespaces(
 	ctx context.Context,
 	request *p.InternalListNamespacesRequest,
 ) (*p.InternalListNamespacesResponse, error) {
-	query := m.session.Query(templateListNamespaceQueryV2, constNamespacePartition).WithContext(ctx)
+	query := m.session.Query(templateListNamespaceQueryV2, constNamespacePartition).WithContext(ctx).Idempotent(true)
 	pageSize := request.PageSize
 	nextPageToken := request.NextPageToken
 	response := &p.InternalListNamespacesResponse{}
@@ -418,7 +418,7 @@ func (m *MetadataStore) DeleteNamespace(
 	request *p.DeleteNamespaceRequest,
 ) error {
 	var name string
-	query := m.session.Query(templateGetNamespaceQuery, request.ID).WithContext(ctx)
+	query := m.session.Query(templateGetNamespaceQuery, request.ID).WithContext(ctx).Idempotent(true)
 	err := query.Scan(&name)
 	if err != nil {
 		if gocql.IsNotFoundError(err) {
@@ -439,7 +439,7 @@ func (m *MetadataStore) DeleteNamespaceByName(
 	request *p.DeleteNamespaceByNameRequest,
 ) error {
 	var ID []byte
-	query := m.session.Query(templateGetNamespaceByNameQueryV2, constNamespacePartition, request.Name).WithContext(ctx)
+	query := m.session.Query(templateGetNamespaceByNameQueryV2, constNamespacePartition, request.Name).WithContext(ctx).Idempotent(true)
 	err := query.Scan(&ID, nil, nil, nil, nil, nil)
 	if err != nil {
 		if gocql.IsNotFoundError(err) {
@@ -454,7 +454,7 @@ func (m *MetadataStore) GetMetadata(
 	ctx context.Context,
 ) (*p.GetMetadataResponse, error) {
 	var notificationVersion int64
-	query := m.session.Query(templateGetMetadataQueryV2, constNamespacePartition, namespaceMetadataRecordName).WithContext(ctx)
+	query := m.session.Query(templateGetMetadataQueryV2, constNamespacePartition, namespaceMetadataRecordName).WithContext(ctx).Idempotent(true)
 	err := query.Scan(&notificationVersion)
 	if err != nil {
 		if gocql.IsNotFoundError(err) {
@@ -486,12 +486,12 @@ func (m *MetadataStore) updateMetadataBatch(
 }
 
 func (m *MetadataStore) deleteNamespace(ctx context.Context, name string, ID []byte) error {
-	query := m.session.Query(templateDeleteNamespaceByNameQueryV2, constNamespacePartition, name).WithContext(ctx)
+	query := m.session.Query(templateDeleteNamespaceByNameQueryV2, constNamespacePartition, name).WithContext(ctx).Idempotent(true)
 	if err := query.Exec(); err != nil {
 		return gocql.ConvertError("DeleteNamespaceByName", err)
 	}
 
-	query = m.session.Query(templateDeleteNamespaceQuery, ID).WithContext(ctx)
+	query = m.session.Query(templateDeleteNamespaceQuery, ID).WithContext(ctx).Idempotent(true)
 	if err := query.Exec(); err != nil {
 		return gocql.ConvertError("DeleteNamespace", err)
 	}

common/persistence/cassandra/mutable_state_store.go

@@ -504,7 +504,7 @@ func (d *MutableStateStore) GetWorkflowExecution(
 		request.RunID,
 		defaultVisibilityTimestamp,
 		rowTypeExecutionTaskID,
-	).WithContext(ctx)
+	).WithContext(ctx).Idempotent(true)
 
 	result := make(map[string]any)
 	if err := query.MapScan(result); err != nil {
@@ -930,7 +930,7 @@ func (d *MutableStateStore) DeleteWorkflowExecution(
 		request.RunID,
 		defaultVisibilityTimestamp,
 		rowTypeExecutionTaskID,
-	).WithContext(ctx)
+	).WithContext(ctx).Idempotent(true)
 
 	err := query.Exec()
 	return gocql.ConvertError("DeleteWorkflowExecution", err)
@@ -967,7 +967,7 @@ func (d *MutableStateStore) GetCurrentExecution(
 		d.getCurrentRecordRunID(request.ArchetypeID),
 		defaultVisibilityTimestamp,
 		rowTypeExecutionTaskID,
-	).WithContext(ctx)
+	).WithContext(ctx).Idempotent(true)
 
 	result := make(map[string]any)
 	if err := query.MapScan(result); err != nil {
@@ -1056,7 +1056,7 @@ func (d *MutableStateStore) ListConcreteExecutions(
 		templateListWorkflowExecutionQuery,
 		request.ShardID,
 		rowTypeExecution,
-	).WithContext(ctx)
+	).WithContext(ctx).Idempotent(true)
 	iter := query.PageSize(request.PageSize).PageState(request.PageToken).Iter()
 
 	response := &p.InternalListConcreteExecutionsResponse{}