forward: add --dry-run

Author: ethanyzhang

cmd/forward.go

@@ -46,6 +46,7 @@ func init() {
 	forward.PrestoFlagsArray.Install(forwardCmd)
 	_ = forwardCmd.Flags().MarkHidden("trino")
 	wd, _ := os.Getwd()
+	forwardCmd.Flags().BoolVarP(&forward.DryRun, "dry-run", "", false, "Turning on dry run will only show the queries but not sending them to the target server.")
 	forwardCmd.Flags().StringVarP(&forward.OutputPath, "output-path", "o", wd, "Output directory path")
 	forwardCmd.Flags().StringVarP(&forward.RunName, "name", "n", fmt.Sprintf("forward_%s", time.Now().Format(utils.DirectoryNameTimeFormat)), `Assign a name to this run. (default: "forward_<current time>")`)
 	forwardCmd.Flags().DurationVarP(&forward.PollInterval, "poll-interval", "i", time.Second*5, "Interval between polls to the source cluster")

cmd/forward/main.go

@@ -18,6 +18,7 @@ import (
 )
 
 var (
+	DryRun                bool
 	PrestoFlagsArray      utils.PrestoFlagsArray
 	OutputPath            string
 	RunName               string
@@ -172,7 +173,22 @@ func forwardQuery(ctx context.Context, queryState *presto.QueryStateInfo, client
 			queryInfo.Query = replacedQuery
 		}
 	}
+	if schema := queryInfo.Session.Schema; schema != nil {
+		if mappedSchema, exists := schemaMappings[*schema]; exists {
+			queryInfo.Session.Schema = &mappedSchema
+			log.Info().Str("source_query_id", queryInfo.QueryId).
+				Msgf("schema replaced %s -> %s", *schema, mappedSchema)
+		}
+	}
 	SessionPropertyHeader := clients[0].GenerateSessionParamsHeaderValue(queryInfo.Session.CollectSessionProperties())
+	if DryRun {
+		logEntry := log.Info().Str("query", queryInfo.Query)
+		if queryInfo.Session.Schema != nil {
+			logEntry = logEntry.Str("schema", *queryInfo.Session.Schema)
+		}
+		logEntry.Msg("query not sent in dry-run mode")
+		return
+	}
 	successful, failed := atomic.Uint32{}, atomic.Uint32{}
 	forwardedQueries := sync.WaitGroup{}
 	for i := 1; i < len(clients); i++ {
@@ -184,13 +200,7 @@ func forwardQuery(ctx context.Context, queryState *presto.QueryStateInfo, client
 					req.Header.Set(presto.CatalogHeader, *queryInfo.Session.Catalog)
 				}
 				if queryInfo.Session.Schema != nil {
-					if mappedSchema, exists := schemaMappings[*queryInfo.Session.Schema]; exists {
-						req.Header.Set(presto.SchemaHeader, mappedSchema)
-						log.Info().Str("source_query_id", queryInfo.QueryId).
-							Msgf("schema replaced %s -> %s", *queryInfo.Session.Schema, mappedSchema)
-					} else {
-						req.Header.Set(presto.SchemaHeader, *queryInfo.Session.Schema)
-					}
+					req.Header.Set(presto.SchemaHeader, *queryInfo.Session.Schema)
 				}
 				req.Header.Set(presto.SessionHeader, SessionPropertyHeader)
 				req.Header.Set(presto.SourceHeader, queryInfo.QueryId)