Merge branch 'master' into coverdatastore

uber · Mar 14, 2024 · 2907992 · 2907992
2 parents 6e4b033 + 6c1ca00
commit 2907992
Showing 1 changed file with 26 additions and 13 deletions.
diff --git a/common/persistence/pinotVisibilityTripleManager.go b/common/persistence/pinotVisibilityTripleManager.go
@@ -329,11 +329,13 @@ type userParameters struct {
 
 // For Pinot Migration uses. It will be a temporary usage
 // logUserQueryParameters will log user queries' parameters so that a comparator workflow can consume
-func (v *pinotVisibilityTripleManager) logUserQueryParameters(userParam userParameters, domain string) {
+func (v *pinotVisibilityTripleManager) logUserQueryParameters(userParam userParameters, domain string, override bool) {
  // Don't log if it is not enabled
- if !v.logCustomerQueryParameter(domain) {
+ // don't log if it is a call from Pinot Response Comparator workflow
+ if !v.logCustomerQueryParameter(domain) || override {
  return
  }
+
  randNum := rand.Intn(10)
  if randNum != 5 { // Intentionally to have 1/10 chance to log custom query parameters
  return
@@ -362,13 +364,14 @@ func (v *pinotVisibilityTripleManager) ListOpenWorkflowExecutions(
  ctx context.Context,
  request *ListWorkflowExecutionsRequest,
 ) (*ListWorkflowExecutionsResponse, error) {
+ override := ctx.Value(ContextKey)
  v.logUserQueryParameters(userParameters{
  operation: string(Operation.LIST),
  domainName: request.Domain,
  closeStatus: -1, // is open. Will have --open flag in comparator workflow
  earliestTime: request.EarliestTime,
  latestTime: request.LatestTime,
- }, request.Domain)
+ }, request.Domain, override != nil)
 
  manager := v.chooseVisibilityManagerForRead(ctx, request.Domain)
  return manager.ListOpenWorkflowExecutions(ctx, request)
@@ -378,13 +381,14 @@ func (v *pinotVisibilityTripleManager) ListClosedWorkflowExecutions(
  ctx context.Context,
  request *ListWorkflowExecutionsRequest,
 ) (*ListWorkflowExecutionsResponse, error) {
+ override := ctx.Value(ContextKey)
  v.logUserQueryParameters(userParameters{
  operation: string(Operation.LIST),
  domainName: request.Domain,
  closeStatus: 6, // 6 means not set closeStatus.
  earliestTime: request.EarliestTime,
  latestTime: request.LatestTime,
- }, request.Domain)
+ }, request.Domain, override != nil)
  manager := v.chooseVisibilityManagerForRead(ctx, request.Domain)
  return manager.ListClosedWorkflowExecutions(ctx, request)
 }
@@ -393,14 +397,15 @@ func (v *pinotVisibilityTripleManager) ListOpenWorkflowExecutionsByType(
  ctx context.Context,
  request *ListWorkflowExecutionsByTypeRequest,
 ) (*ListWorkflowExecutionsResponse, error) {
+ override := ctx.Value(ContextKey)
  v.logUserQueryParameters(userParameters{
  operation: string(Operation.LIST),
  domainName: request.Domain,
  workflowType: request.WorkflowTypeName,
  closeStatus: -1, // is open. Will have --open flag in comparator workflow
  earliestTime: request.EarliestTime,
  latestTime: request.LatestTime,
- }, request.Domain)
+ }, request.Domain, override != nil)
  manager := v.chooseVisibilityManagerForRead(ctx, request.Domain)
  return manager.ListOpenWorkflowExecutionsByType(ctx, request)
 }
@@ -409,14 +414,15 @@ func (v *pinotVisibilityTripleManager) ListClosedWorkflowExecutionsByType(
  ctx context.Context,
  request *ListWorkflowExecutionsByTypeRequest,
 ) (*ListWorkflowExecutionsResponse, error) {
+ override := ctx.Value(ContextKey)
  v.logUserQueryParameters(userParameters{
  operation: string(Operation.LIST),
  domainName: request.Domain,
  workflowType: request.WorkflowTypeName,
  closeStatus: 6, // 6 means not set closeStatus.
  earliestTime: request.EarliestTime,
  latestTime: request.LatestTime,
- }, request.Domain)
+ }, request.Domain, override != nil)
  manager := v.chooseVisibilityManagerForRead(ctx, request.Domain)
  return manager.ListClosedWorkflowExecutionsByType(ctx, request)
 }
@@ -425,14 +431,15 @@ func (v *pinotVisibilityTripleManager) ListOpenWorkflowExecutionsByWorkflowID(
  ctx context.Context,
  request *ListWorkflowExecutionsByWorkflowIDRequest,
 ) (*ListWorkflowExecutionsResponse, error) {
+ override := ctx.Value(ContextKey)
  v.logUserQueryParameters(userParameters{
  operation: string(Operation.LIST),
  domainName: request.Domain,
  workflowID: request.WorkflowID,
  closeStatus: -1,
  earliestTime: request.EarliestTime,
  latestTime: request.LatestTime,
- }, request.Domain)
+ }, request.Domain, override != nil)
  manager := v.chooseVisibilityManagerForRead(ctx, request.Domain)
  return manager.ListOpenWorkflowExecutionsByWorkflowID(ctx, request)
 }
@@ -441,14 +448,15 @@ func (v *pinotVisibilityTripleManager) ListClosedWorkflowExecutionsByWorkflowID(
  ctx context.Context,
  request *ListWorkflowExecutionsByWorkflowIDRequest,
 ) (*ListWorkflowExecutionsResponse, error) {
+ override := ctx.Value(ContextKey)
  v.logUserQueryParameters(userParameters{
  operation: string(Operation.LIST),
  domainName: request.Domain,
  workflowID: request.WorkflowID,
  closeStatus: 6, // 6 means not set closeStatus.
  earliestTime: request.EarliestTime,
  latestTime: request.LatestTime,
- }, request.Domain)
+ }, request.Domain, override != nil)
  manager := v.chooseVisibilityManagerForRead(ctx, request.Domain)
  return manager.ListClosedWorkflowExecutionsByWorkflowID(ctx, request)
 }
@@ -457,13 +465,14 @@ func (v *pinotVisibilityTripleManager) ListClosedWorkflowExecutionsByStatus(
  ctx context.Context,
  request *ListClosedWorkflowExecutionsByStatusRequest,
 ) (*ListWorkflowExecutionsResponse, error) {
+ override := ctx.Value(ContextKey)
  v.logUserQueryParameters(userParameters{
  operation: string(Operation.LIST),
  domainName: request.Domain,
  closeStatus: int(request.Status),
  earliestTime: request.EarliestTime,
  latestTime: request.LatestTime,
- }, request.Domain)
+ }, request.Domain, override != nil)
  manager := v.chooseVisibilityManagerForRead(ctx, request.Domain)
  return manager.ListClosedWorkflowExecutionsByStatus(ctx, request)
 }
@@ -475,13 +484,14 @@ func (v *pinotVisibilityTripleManager) GetClosedWorkflowExecution(
  earlistTime := int64(0) // this is to get all closed workflow execution
  latestTime := time.Now().UnixNano()
 
+ override := ctx.Value(ContextKey)
  v.logUserQueryParameters(userParameters{
  operation: string(Operation.LIST),
  domainName: request.Domain,
  closeStatus: 6, // 6 means not set closeStatus.
  earliestTime: earlistTime,
  latestTime: latestTime,
- }, request.Domain)
+ }, request.Domain, override != nil)
  manager := v.chooseVisibilityManagerForRead(ctx, request.Domain)
  return manager.GetClosedWorkflowExecution(ctx, request)
 }
@@ -490,14 +500,15 @@ func (v *pinotVisibilityTripleManager) ListWorkflowExecutions(
  ctx context.Context,
  request *ListWorkflowExecutionsByQueryRequest,
 ) (*ListWorkflowExecutionsResponse, error) {
+ override := ctx.Value(ContextKey)
  v.logUserQueryParameters(userParameters{
  operation: string(Operation.LIST),
  domainName: request.Domain,
  closeStatus: 6, // 6 means not set closeStatus.
  customQuery: request.Query,
  earliestTime: -1,
  latestTime: -1,
- }, request.Domain)
+ }, request.Domain, override != nil)
  manager := v.chooseVisibilityManagerForRead(ctx, request.Domain)
  return manager.ListWorkflowExecutions(ctx, request)
 }
@@ -506,14 +517,15 @@ func (v *pinotVisibilityTripleManager) ScanWorkflowExecutions(
  ctx context.Context,
  request *ListWorkflowExecutionsByQueryRequest,
 ) (*ListWorkflowExecutionsResponse, error) {
+ override := ctx.Value(ContextKey)
  v.logUserQueryParameters(userParameters{
  operation: string(Operation.LIST),
  domainName: request.Domain,
  closeStatus: 6, // 6 means not set closeStatus.
  customQuery: request.Query,
  earliestTime: -1,
  latestTime: -1,
- }, request.Domain)
+ }, request.Domain, override != nil)
  manager := v.chooseVisibilityManagerForRead(ctx, request.Domain)
  return manager.ScanWorkflowExecutions(ctx, request)
 }
@@ -522,14 +534,15 @@ func (v *pinotVisibilityTripleManager) CountWorkflowExecutions(
  ctx context.Context,
  request *CountWorkflowExecutionsRequest,
 ) (*CountWorkflowExecutionsResponse, error) {
+ override := ctx.Value(ContextKey)
  v.logUserQueryParameters(userParameters{
  operation: string(Operation.COUNT),
  domainName: request.Domain,
  closeStatus: 6, // 6 means not set closeStatus.
  customQuery: request.Query,
  earliestTime: -1,
  latestTime: -1,
- }, request.Domain)
+ }, request.Domain, override != nil)
  manager := v.chooseVisibilityManagerForRead(ctx, request.Domain)
  return manager.CountWorkflowExecutions(ctx, request)
 }