Repository: uber/cadence
Pull request #3072 (2020-02-22 23:38:19)

S3 visiblity store - Add ability to query WorkflowTypeName
Change layout of stored objects on s3

coveralls: 
[![Coverage Status](https://coveralls.io/builds/28928872/badge)](https://coveralls.io/builds/28928872)

Coverage increased (+0.01%) to 67.357% when pulling **a83fc9c9e6d3bcfb31fe6af08ae6fcc570f29199 on jontro:s3-archival-indexing** into **c09186f4c8583d3084315b9fe8a493c2b080d1f6 on uber:master**.


Patch:
diff --git a/common/archiver/s3store/README.md b/common/archiver/s3store/README.md
index fff81097a..a740e0a68 100644
--- a/common/archiver/s3store/README.md
+++ b/common/archiver/s3store/README.md
@@ -35,11 +35,12 @@ The syntax for the query is based on SQL
 
 Supported column names are
 - WorkflowID *String*
+- WorkflowTypeName *String*
 - StartTime *Date*
 - CloseTime *Date*
 - SearchPrecision *String - Day, Hour, Minute, Second*
 
-WorkflowID and SearchPrecision are always required. One of StartTime and CloseTime are required and they are mutually exclusive.
+WorkflowID or WorkflowTypeName is required. If filtering on date use StartTime or CloseTime in combination with SearchPrecision.
 
 Searching for a record will be done in times in the UTC timezone
 
@@ -58,11 +59,15 @@ it will search all records starting from `2020-01-21T00:00:00Z` to `2020-01-21T5
 ## Storage in S3
 Workflow runs are stored in s3 using the following structure
 ```
-s3://<bucket-name>/<domain-id>/<workflow-id>/
-	history/<run-id>
+s3://<bucket-name>/<domain-id>/
+	history/<workflow-id>/<run-id>
 	visibility/
-		startTimeout/2020-01-21T16:16:11Z/<run-id>
-		closeTimeout/2020-01-21T16:16:11Z/<run-id>
+            workflowTypeName/<workflow-type-name>/
+                startTimeout/2020-01-21T16:16:11Z/<run-id>
+                closeTimeout/2020-01-21T16:16:11Z/<run-id>
+            workflowID/<workflow-id>/
+                startTimeout/2020-01-21T16:16:11Z/<run-id>
+                closeTimeout/2020-01-21T16:16:11Z/<run-id>
 ```
 
 ## Using localstack for local development
diff --git a/common/archiver/s3store/queryParser.go b/common/archiver/s3store/queryParser.go
index fcbcd3112..da58817e3 100644
--- a/common/archiver/s3store/queryParser.go
+++ b/common/archiver/s3store/queryParser.go
@@ -42,20 +42,21 @@ type (
 	queryParser struct{}
 
 	parsedQuery struct {
-		workflowID      *string
-		startTime       *int64
-		closeTime       *int64
-		searchPrecision *string
-		emptyResult     bool
+		workflowTypeName *string
+		workflowID       *string
+		startTime        *int64
+		closeTime        *int64
+		searchPrecision  *string
 	}
 )
 
 // All allowed fields for filtering
 const (
-	WorkflowID      = "WorkflowID"
-	StartTime       = "StartTime"
-	CloseTime       = "CloseTime"
-	SearchPrecision = "SearchPrecision"
+	WorkflowTypeName = "WorkflowTypeName"
+	WorkflowID       = "WorkflowID"
+	StartTime        = "StartTime"
+	CloseTime        = "CloseTime"
+	SearchPrecision  = "SearchPrecision"
 )
 
 // Precision specific values
@@ -85,8 +86,11 @@ func (p *queryParser) Parse(query string) (*parsedQuery, error) {
 	if err := p.convertWhereExpr(whereExpr, parsedQuery); err != nil {
 		return nil, err
 	}
-	if parsedQuery.workflowID == nil {
-		return nil, errors.New("workflowID is required in query")
+	if parsedQuery.workflowID == nil && parsedQuery.workflowTypeName == nil {
+		return nil, errors.New("WorkflowID or WorkflowTypeName is required in query")
+	}
+	if parsedQuery.workflowID != nil && parsedQuery.workflowTypeName != nil {
+		return nil, errors.New("only one of WorkflowID or WorkflowTypeName can be specified in a query")
 	}
 	if parsedQuery.closeTime != nil && parsedQuery.startTime != nil {
 		return nil, errors.New("only one of StartTime or CloseTime can be specified in a query")
@@ -143,6 +147,18 @@ func (p *queryParser) convertComparisonExpr(compExpr *sqlparser.ComparisonExpr,
 	valStr := sqlparser.String(valExpr)
 
 	switch colNameStr {
+	case WorkflowTypeName:
+		val, err := extractStringValue(valStr)
+		if err != nil {
+			return err
+		}
+		if op != "=" {
+			return fmt.Errorf("only operation = is support for %s", WorkflowTypeName)
+		}
+		if parsedQuery.workflowTypeName != nil {
+			return fmt.Errorf("can not query %s multiple times", WorkflowTypeName)
+		}
+		parsedQuery.workflowTypeName = common.StringPtr(val)
 	case WorkflowID:
 		val, err := extractStringValue(valStr)
 		if err != nil {
@@ -151,9 +167,8 @@ func (p *queryParser) convertComparisonExpr(compExpr *sqlparser.ComparisonExpr,
 		if op != "=" {
 			return fmt.Errorf("only operation = is support for %s", WorkflowID)
 		}
-		if parsedQuery.workflowID != nil && *parsedQuery.workflowID != val {
-			parsedQuery.emptyResult = true
-			return nil
+		if parsedQuery.workflowID != nil {
+			return fmt.Errorf("can not query %s multiple times", WorkflowID)
 		}
 		parsedQuery.workflowID = common.StringPtr(val)
 	case CloseTime:
diff --git a/common/archiver/s3store/util.go b/common/archiver/s3store/util.go
index 66bd5ea48..bd2c73fc4 100644
--- a/common/archiver/s3store/util.go
+++ b/common/archiver/s3store/util.go
@@ -27,6 +27,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"strings"
+	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/awserr"
@@ -119,16 +120,36 @@ func constructHistoryKey(path, domainID, workflowID, runID string, version int64
 }
 
 func constructHistoryKeyPrefix(path, domainID, workflowID, runID string) string {
-	return fmt.Sprintf("%s/%s", constructCommonKeyPrefix(path, domainID, workflowID, "history"), runID)
+	return strings.TrimLeft(strings.Join([]string{path, domainID, "history", workflowID, runID}, "/"), "/")
 }
 
-func constructVisibilitySearchPrefix(path, domainID, workflowID, indexType string) string {
-	return constructCommonKeyPrefix(path, domainID, workflowID, fmt.Sprintf("visibility/%s", indexType))
+func constructTimeBasedSearchKey(path, domainID, primaryIndexKey, primaryIndexValue, secondaryIndexKey string, timestamp int64, precision string) string {
+	t := time.Unix(0, timestamp).In(time.UTC)
+	var timeFormat = ""
+	switch precision {
+	case PrecisionSecond:
+		timeFormat = ":05"
+		fallthrough
+	case PrecisionMinute:
+		timeFormat = ":04" + timeFormat
+		fallthrough
+	case PrecisionHour:
+		timeFormat = "15" + timeFormat
+		fallthrough
+	case PrecisionDay:
+		timeFormat = "2006-01-02T" + timeFormat
+	}
+
+	return fmt.Sprintf("%s/%s", constructVisibilitySearchPrefix(path, domainID, primaryIndexKey, primaryIndexValue, secondaryIndexKey), t.Format(timeFormat))
+}
+
+func constructTimestampIndex(path, domainID, primaryIndexKey, primaryIndexValue, secondaryIndexKey string, timestamp int64, runID string) string {
+	t := time.Unix(0, timestamp).In(time.UTC)
+	return fmt.Sprintf("%s/%s/%s", constructVisibilitySearchPrefix(path, domainID, primaryIndexKey, primaryIndexValue, secondaryIndexKey), t.Format(time.RFC3339), runID)
 }
 
-// Make sure there are no trailing slashes
-func constructCommonKeyPrefix(path, domainID, workflowID, entryType string) string {
-	return strings.TrimLeft(strings.Join([]string{path, domainID, workflowID, entryType}, "/"), "/")
+func constructVisibilitySearchPrefix(path, domainID, primaryIndexKey, primaryIndexValue, secondaryIndexType string) string {
+	return strings.TrimLeft(strings.Join([]string{path, domainID, "visibility", primaryIndexKey, primaryIndexValue, secondaryIndexType}, "/"), "/")
 }
 
 func ensureContextTimeout(ctx context.Context) (context.Context, context.CancelFunc) {
diff --git a/common/archiver/s3store/visibilityArchiver.go b/common/archiver/s3store/visibilityArchiver.go
index eafb3ae43..7b5962271 100644
--- a/common/archiver/s3store/visibilityArchiver.go
+++ b/common/archiver/s3store/visibilityArchiver.go
@@ -22,8 +22,6 @@ package s3store
 
 import (
 	"context"
-	"fmt"
-	"time"
 
 	"github.com/uber/cadence/common/metrics"
 
@@ -53,12 +51,21 @@ type (
 		nextPageToken []byte
 		parsedQuery   *parsedQuery
 	}
+
+	indexToArchive struct {
+		primaryIndex            string
+		primaryIndexValue       string
+		secondaryIndex          string
+		secondaryIndexTimestamp int64
+	}
 )
 
 const (
-	errEncodeVisibilityRecord = "failed to encode visibility record"
-	indexKeyStartTimeout      = "startTimeout"
-	indexKeyCloseTimeout      = "closeTimeout"
+	errEncodeVisibilityRecord       = "failed to encode visibility record"
+	secondaryIndexKeyStartTimeout   = "startTimeout"
+	secondaryIndexKeyCloseTimeout   = "closeTimeout"
+	primaryIndexKeyWorkflowTypeName = "workflowTypeName"
+	primaryIndexKeyWorkflowID       = "workflowID"
 )
 
 // NewVisibilityArchiver creates a new archiver.VisibilityArchiver based on s3
@@ -98,13 +105,17 @@ func (v *visibilityArchiver) Archive(
 	scope := v.container.MetricsClient.Scope(metrics.VisibilityArchiverScope, metrics.DomainTag(request.DomainName))
 	featureCatalog := archiver.GetFeatureCatalog(opts...)
 	sw := scope.StartTimer(metrics.CadenceLatency)
+	logger := archiver.TagLoggerWithArchiveVisibilityRequestAndURI(v.container.Logger, request, URI.String())
+	archiveFailReason := ""
 	defer func() {
 		sw.Stop()
 		if err != nil {
 			if isRetryableError(err) {
 				scope.IncCounter(metrics.VisibilityArchiverArchiveTransientErrorCount)
+				logger.Error(archiver.ArchiveTransientErrorMsg, tag.ArchivalArchiveFailReason(archiveFailReason), tag.Error(err))
 			} else {
 				scope.IncCounter(metrics.VisibilityArchiverArchiveNonRetryableErrorCount)
+				logger.Error(archiver.ArchiveNonRetriableErrorMsg, tag.ArchivalArchiveFailReason(archiveFailReason), tag.Error(err))
 				if featureCatalog.NonRetriableError != nil {
 					err = featureCatalog.NonRetriableError()
 				}
@@ -112,62 +123,42 @@ func (v *visibilityArchiver) Archive(
 		}
 	}()
 
-	logger := archiver.TagLoggerWithArchiveVisibilityRequestAndURI(v.container.Logger, request, URI.String())
-
 	if err := softValidateURI(URI); err != nil {
-		logger.Error(archiver.ArchiveNonRetriableErrorMsg, tag.ArchivalArchiveFailReason(archiver.ErrReasonInvalidURI), tag.Error(err))
+		archiveFailReason = archiver.ErrReasonInvalidURI
 		return err
 	}
 
 	if err := archiver.ValidateVisibilityArchivalRequest(request); err != nil {
-		logger.Error(archiver.ArchiveNonRetriableErrorMsg, tag.ArchivalArchiveFailReason(archiver.ErrReasonInvalidArchiveRequest), tag.Error(err))
+		archiveFailReason = archiver.ErrReasonInvalidArchiveRequest
 		return err
 	}
 
 	encodedVisibilityRecord, err := encode(request)
 	if err != nil {
-		logger.Error(archiver.ArchiveNonRetriableErrorMsg, tag.ArchivalArchiveFailReason(errEncodeVisibilityRecord), tag.Error(err))
+		archiveFailReason = errEncodeVisibilityRecord
 		return err
 	}
-
+	indexes := createIndexesToArchive(request)
 	// Upload archive to all indexes
-	key := constructTimestampIndex(URI.Path(), request.DomainID, request.WorkflowID, indexKeyCloseTimeout, request.CloseTimestamp, request.RunID)
-	if err := upload(ctx, v.s3cli, URI, key, encodedVisibilityRecord); err != nil {
-		logger.Error(archiver.ArchiveNonRetriableErrorMsg, tag.ArchivalArchiveFailReason(errWriteKey), tag.Error(err))
-		return err
-	}
-	key = constructTimestampIndex(URI.Path(), request.DomainID, request.WorkflowID, indexKeyStartTimeout, request.StartTimestamp, request.RunID)
-	if err := upload(ctx, v.s3cli, URI, key, encodedVisibilityRecord); err != nil {
-		logger.Error(archiver.ArchiveNonRetriableErrorMsg, tag.ArchivalArchiveFailReason(errWriteKey), tag.Error(err))
-		return err
+	for _, element := range indexes {
+		key := constructTimestampIndex(URI.Path(), request.DomainID, element.primaryIndex, element.primaryIndexValue, element.secondaryIndex, element.secondaryIndexTimestamp, request.RunID)
+		if err := upload(ctx, v.s3cli, URI, key, encodedVisibilityRecord); err != nil {
+			archiveFailReason = errWriteKey
+			return err
+		}
 	}
 	scope.IncCounter(metrics.VisibilityArchiveSuccessCount)
 	return nil
 }
-
-func constructTimeBasedSearchKey(path, domainID, workflowID, indexKey string, timestamp int64, precision string) string {
-	t := time.Unix(0, timestamp).In(time.UTC)
-	var timeFormat = ""
-	switch precision {
-	case PrecisionSecond:
-		timeFormat = ":05"
-		fallthrough
-	case PrecisionMinute:
-		timeFormat = ":04" + timeFormat
-		fallthrough
-	case PrecisionHour:
-		timeFormat = "15" + timeFormat
-		fallthrough
-	case PrecisionDay:
-		timeFormat = "2006-01-02T" + timeFormat
+func createIndexesToArchive(request *archiver.ArchiveVisibilityRequest) []indexToArchive {
+	return []indexToArchive{
+		{primaryIndexKeyWorkflowTypeName, request.WorkflowTypeName, secondaryIndexKeyCloseTimeout, request.CloseTimestamp},
+		{primaryIndexKeyWorkflowTypeName, request.WorkflowTypeName, secondaryIndexKeyStartTimeout, request.StartTimestamp},
+		{primaryIndexKeyWorkflowID, request.WorkflowID, secondaryIndexKeyCloseTimeout, request.CloseTimestamp},
+		{primaryIndexKeyWorkflowID, request.WorkflowID, secondaryIndexKeyStartTimeout, request.StartTimestamp},
 	}
-
-	return fmt.Sprintf("%s/%s", constructVisibilitySearchPrefix(path, domainID, workflowID, indexKey), t.Format(timeFormat))
-}
-func constructTimestampIndex(path, domainID, workflowID, indexKey string, timestamp int64, runID string) string {
-	t := time.Unix(0, timestamp).In(time.UTC)
-	return fmt.Sprintf("%s/%s/%s", constructVisibilitySearchPrefix(path, domainID, workflowID, indexKey), t.Format(time.RFC3339), runID)
 }
+
 func (v *visibilityArchiver) Query(
 	ctx context.Context,
 	URI archiver.URI,
@@ -185,9 +176,6 @@ func (v *visibilityArchiver) Query(
 	if err != nil {
 		return nil, &shared.BadRequestError{Message: err.Error()}
 	}
-	if parsedQuery.emptyResult {
-		return &archiver.QueryVisibilityResponse{}, nil
-	}
 
 	return v.query(ctx, URI, &queryVisibilityRequest{
 		domainID:      request.DomainID,
@@ -208,12 +196,18 @@ func (v *visibilityArchiver) query(
 	if request.nextPageToken != nil {
 		token = deserializeQueryVisibilityToken(request.nextPageToken)
 	}
-	var prefix = constructVisibilitySearchPrefix(URI.Path(), request.domainID, *request.parsedQuery.workflowID, indexKeyCloseTimeout) + "/"
+	primaryIndex := primaryIndexKeyWorkflowTypeName
+	primaryIndexValue := request.parsedQuery.workflowTypeName
+	if request.parsedQuery.workflowID != nil {
+		primaryIndex = primaryIndexKeyWorkflowID
+		primaryIndexValue = request.parsedQuery.workflowID
+	}
+	var prefix = constructVisibilitySearchPrefix(URI.Path(), request.domainID, primaryIndex, *primaryIndexValue, secondaryIndexKeyCloseTimeout) + "/"
 	if request.parsedQuery.closeTime != nil {
-		prefix = constructTimeBasedSearchKey(URI.Path(), request.domainID, *request.parsedQuery.workflowID, indexKeyCloseTimeout, *request.parsedQuery.closeTime, *request.parsedQuery.searchPrecision)
+		prefix = constructTimeBasedSearchKey(URI.Path(), request.domainID, primaryIndex, *primaryIndexValue, secondaryIndexKeyCloseTimeout, *request.parsedQuery.closeTime, *request.parsedQuery.searchPrecision)
 	}
 	if request.parsedQuery.startTime != nil {
-		prefix = constructTimeBasedSearchKey(URI.Path(), request.domainID, *request.parsedQuery.workflowID, indexKeyStartTimeout, *request.parsedQuery.startTime, *request.parsedQuery.searchPrecision)
+		prefix = constructTimeBasedSearchKey(URI.Path(), request.domainID, primaryIndex, *primaryIndexValue, secondaryIndexKeyStartTimeout, *request.parsedQuery.startTime, *request.parsedQuery.searchPrecision)
 	}
 
 	results, err := v.s3cli.ListObjectsV2WithContext(ctx, &s3.ListObjectsV2Input{


Test patch:
diff --git a/common/archiver/s3store/queryParser_test.go b/common/archiver/s3store/queryParser_test.go
index e4e0526de..7530c1f92 100644
--- a/common/archiver/s3store/queryParser_test.go
+++ b/common/archiver/s3store/queryParser_test.go
@@ -45,7 +45,7 @@ func (s *queryParserSuite) SetupTest() {
 	s.parser = NewQueryParser()
 }
 
-func (s *queryParserSuite) TestParseWorkflowID() {
+func (s *queryParserSuite) TestParseWorkflowIDAndWorkflowTypeName() {
 	testCases := []struct {
 		query       string
 		expectErr   bool
@@ -59,12 +59,20 @@ func (s *queryParserSuite) TestParseWorkflowID() {
 			},
 		},
 		{
-			query:     "WorkflowID = \"random workflowID\" and WorkflowID = \"random workflowID\"",
+			query:     "WorkflowTypeName = \"random workflowTypeName\"",
 			expectErr: false,
 			parsedQuery: &parsedQuery{
-				workflowID: common.StringPtr("random workflowID"),
+				workflowTypeName: common.StringPtr("random workflowTypeName"),
 			},
 		},
+		{
+			query:     "WorkflowID = \"random workflowID\" and WorkflowTypeName = \"random workflowTypeName\"",
+			expectErr: true,
+		},
+		{
+			query:     "WorkflowID = \"random workflowID\" and WorkflowID = \"random workflowID\"",
+			expectErr: true,
+		},
 		{
 			query:     "RunID = \"random runID\"",
 			expectErr: true,
@@ -112,10 +120,9 @@ func (s *queryParserSuite) TestParseWorkflowID() {
 			continue
 		}
 		s.NoError(err)
-		s.Equal(tc.parsedQuery.emptyResult, parsedQuery.emptyResult)
-		if !tc.parsedQuery.emptyResult {
-			s.Equal(tc.parsedQuery.workflowID, parsedQuery.workflowID)
-		}
+		s.Equal(tc.parsedQuery.workflowID, parsedQuery.workflowID)
+		s.Equal(tc.parsedQuery.workflowTypeName, parsedQuery.workflowTypeName)
+
 	}
 }
 
@@ -214,10 +221,8 @@ func (s *queryParserSuite) TestParseCloseTime() {
 			continue
 		}
 		s.NoError(err)
-		s.Equal(tc.parsedQuery.emptyResult, parsedQuery.emptyResult)
-		if !tc.parsedQuery.emptyResult {
-			s.Equal(tc.parsedQuery.closeTime, parsedQuery.closeTime)
-		}
+		s.Equal(tc.parsedQuery.closeTime, parsedQuery.closeTime)
+
 	}
 }
 
@@ -260,9 +265,6 @@ func (s *queryParserSuite) TestParseStartTime() {
 			continue
 		}
 		s.NoError(err)
-		s.Equal(tc.parsedQuery.emptyResult, parsedQuery.emptyResult)
-		if !tc.parsedQuery.emptyResult {
-			s.Equal(tc.parsedQuery.closeTime, parsedQuery.closeTime)
-		}
+		s.Equal(tc.parsedQuery.closeTime, parsedQuery.closeTime)
 	}
 }
diff --git a/common/archiver/s3store/visibilityArchiver_test.go b/common/archiver/s3store/visibilityArchiver_test.go
index 49b975d9f..e0afebe62 100644
--- a/common/archiver/s3store/visibilityArchiver_test.go
+++ b/common/archiver/s3store/visibilityArchiver_test.go
@@ -218,7 +218,7 @@ func (s *visibilityArchiverSuite) TestArchive_Success() {
 	err = visibilityArchiver.Archive(context.Background(), URI, request)
 	s.NoError(err)
 
-	expectedKey := constructTimestampIndex(URI.Path(), testDomainID, testWorkflowID, indexKeyCloseTimeout, closeTimestamp.UnixNano(), testRunID)
+	expectedKey := constructTimestampIndex(URI.Path(), testDomainID, primaryIndexKeyWorkflowID, testWorkflowID, secondaryIndexKeyCloseTimeout, closeTimestamp.UnixNano(), testRunID)
 	data, err := download(context.Background(), visibilityArchiver.s3cli, URI, expectedKey)
 	s.NoError(err, expectedKey)
 
@@ -421,7 +421,7 @@ func (s *visibilityArchiverSuite) TestArchiveAndQueryPrecisions() {
 		},
 	}
 	visibilityArchiver := s.newTestVisibilityArchiver()
-	URI, err := archiver.NewURI(testBucketURI + "/archive-and-query")
+	URI, err := archiver.NewURI(testBucketURI + "/archive-and-query-precision")
 	s.NoError(err)
 
 	for i, testData := range precisionTests {
@@ -472,17 +472,36 @@ func (s *visibilityArchiverSuite) TestArchiveAndQueryPrecisions() {
 		s.NoError(err)
 		s.NotNil(response)
 		s.Len(response.Executions, 2, "Iteration ", i)
+
+		mockParser = NewMockQueryParser(s.controller)
+		mockParser.EXPECT().Parse(gomock.Any()).Return(&parsedQuery{
+			closeTime:        common.Int64Ptr((testData.day+30)*int64(time.Hour)*24 + testData.hour*int64(time.Hour) + testData.minute*int64(time.Minute) + testData.second*int64(time.Second)),
+			searchPrecision:  common.StringPtr(testData.precision),
+			workflowTypeName: common.StringPtr(testWorkflowTypeName),
+		}, nil).AnyTimes()
+		visibilityArchiver.queryParser = mockParser
+
+		response, err = visibilityArchiver.Query(context.Background(), URI, request)
+		s.NoError(err)
+		s.NotNil(response)
+		s.Len(response.Executions, 2, "Iteration ", i)
+
+		mockParser = NewMockQueryParser(s.controller)
+		mockParser.EXPECT().Parse(gomock.Any()).Return(&parsedQuery{
+			startTime:        common.Int64Ptr((testData.day)*int64(time.Hour)*24 + testData.hour*int64(time.Hour) + testData.minute*int64(time.Minute) + testData.second*int64(time.Second)),
+			searchPrecision:  common.StringPtr(testData.precision),
+			workflowTypeName: common.StringPtr(testWorkflowTypeName),
+		}, nil).AnyTimes()
+		visibilityArchiver.queryParser = mockParser
+
+		response, err = visibilityArchiver.Query(context.Background(), URI, request)
+		s.NoError(err)
+		s.NotNil(response)
+		s.Len(response.Executions, 2, "Iteration ", i)
 	}
 }
 func (s *visibilityArchiverSuite) TestArchiveAndQuery() {
 	visibilityArchiver := s.newTestVisibilityArchiver()
-	mockParser := NewMockQueryParser(s.controller)
-	mockParser.EXPECT().Parse(gomock.Any()).Return(&parsedQuery{
-		closeTime:       common.Int64Ptr(int64(1 * time.Hour)),
-		searchPrecision: common.StringPtr(PrecisionHour),
-		workflowID:      common.StringPtr(testWorkflowID),
-	}, nil).AnyTimes()
-	visibilityArchiver.queryParser = mockParser
 	URI, err := archiver.NewURI(testBucketURI + "/archive-and-query")
 	s.NoError(err)
 	for _, record := range s.visibilityRecords {
@@ -490,6 +509,11 @@ func (s *visibilityArchiverSuite) TestArchiveAndQuery() {
 		s.NoError(err)
 	}
 
+	mockParser := NewMockQueryParser(s.controller)
+	mockParser.EXPECT().Parse(gomock.Any()).Return(&parsedQuery{
+		workflowID: common.StringPtr(testWorkflowID),
+	}, nil).AnyTimes()
+	visibilityArchiver.queryParser = mockParser
 	request := &archiver.QueryVisibilityRequest{
 		DomainID: testDomainID,
 		PageSize: 1,
@@ -505,9 +529,35 @@ func (s *visibilityArchiverSuite) TestArchiveAndQuery() {
 		request.NextPageToken = response.NextPageToken
 		first = false
 	}
-	s.Len(executions, 2)
+	s.Len(executions, 3)
+	s.Equal(convertToExecutionInfo(s.visibilityRecords[0]), executions[0])
+	s.Equal(convertToExecutionInfo(s.visibilityRecords[1]), executions[1])
+	s.Equal(convertToExecutionInfo(s.visibilityRecords[2]), executions[2])
+
+	mockParser = NewMockQueryParser(s.controller)
+	mockParser.EXPECT().Parse(gomock.Any()).Return(&parsedQuery{
+		workflowTypeName: common.StringPtr(testWorkflowTypeName),
+	}, nil).AnyTimes()
+	visibilityArchiver.queryParser = mockParser
+	request = &archiver.QueryVisibilityRequest{
+		DomainID: testDomainID,
+		PageSize: 1,
+		Query:    "parsed by mockParser",
+	}
+	executions = []*shared.WorkflowExecutionInfo{}
+	first = true
+	for first || request.NextPageToken != nil {
+		response, err := visibilityArchiver.Query(context.Background(), URI, request)
+		s.NoError(err)
+		s.NotNil(response)
+		executions = append(executions, response.Executions...)
+		request.NextPageToken = response.NextPageToken
+		first = false
+	}
+	s.Len(executions, 3)
 	s.Equal(convertToExecutionInfo(s.visibilityRecords[0]), executions[0])
 	s.Equal(convertToExecutionInfo(s.visibilityRecords[1]), executions[1])
+	s.Equal(convertToExecutionInfo(s.visibilityRecords[2]), executions[2])
 }
 
 func (s *visibilityArchiverSuite) setupVisibilityDirectory() {