feat: add transform function to async destination manager

rudderlabs · Apr 16, 2024 · 95e3e1d · 95e3e1d
1 parent c41ae57
commit 95e3e1d
Show file tree

Hide file tree

Showing 7 changed files with 49 additions and 19 deletions.
diff --git a/router/batchrouter/asyncdestinationmanager/bing-ads/bulk_uploader.go b/router/batchrouter/asyncdestinationmanager/bing-ads/bulk_uploader.go
@@ -13,6 +13,7 @@ import (
 	"github.com/rudderlabs/rudder-go-kit/bytesize"
 	"github.com/rudderlabs/rudder-go-kit/logger"
 	"github.com/rudderlabs/rudder-go-kit/stats"
+	"github.com/rudderlabs/rudder-server/jobsdb"
 	"github.com/rudderlabs/rudder-server/router/batchrouter/asyncdestinationmanager/common"
 	router_utils "github.com/rudderlabs/rudder-server/router/utils"
 	"github.com/rudderlabs/rudder-server/utils/misc"
@@ -29,6 +30,10 @@ func NewBingAdsBulkUploader(destName string, service bingads.BulkServiceI, clien
 	}
 }
 
+func (b *BingAdsBulkUploader) Transform(job *jobsdb.JobT) (string, error) {
+	return common.GetMarshalledData(job), nil
+}
+
 /*
 This function create at most 3 zip files from the text file created by the batchrouter
 It takes the text file path as input and returns the zip file path

diff --git a/router/batchrouter/asyncdestinationmanager/common/common.go b/router/batchrouter/asyncdestinationmanager/common/common.go
@@ -17,6 +17,7 @@ type AsyncDestinationManager interface {
 	Upload(asyncDestStruct *AsyncDestinationStruct) AsyncUploadOutput
 	Poll(pollInput AsyncPoll) PollStatusResponse
 	GetUploadStats(UploadStatsInput GetUploadStatsInput) GetUploadStatsResponse
+	Transform(job *jobsdb.JobT) (string, error)
 }
 
 var AsyncDestinations = []string{"MARKETO_BULK_UPLOAD", "BING_ADS", "ELOQUA"}
@@ -141,6 +142,22 @@ func GetTransformedData(payload stdjson.RawMessage) string {
 	return gjson.Get(string(payload), "body.JSON").String()
 }
 
+func GetMarshalledData(job *jobsdb.JobT) string {
+	payload := GetTransformedData(job.EventPayload)
+	var asyncJob AsyncJob
+	err := stdjson.Unmarshal([]byte(payload), &asyncJob.Message)
+	if err != nil {
+		panic("Unmarshalling Transformer Response Failed")
+	}
+	asyncJob.Metadata = make(map[string]interface{})
+	asyncJob.Metadata["job_id"] = job.JobID
+	responsePayload, err := stdjson.Marshal(asyncJob)
+	if err != nil {
+		panic("Marshalling Response Payload Failed")
+	}
+	return string(responsePayload)
+}
+
 func GetBatchRouterConfigInt64(key, destType string, defaultValue int64) int64 {
 	destOverrideFound := config.IsSet("BatchRouter." + destType + "." + key)
 	if destOverrideFound {

diff --git a/router/batchrouter/asyncdestinationmanager/common/manager.go b/router/batchrouter/asyncdestinationmanager/common/manager.go
@@ -1,11 +1,18 @@
 package common
 
 import (
+	"errors"
 	"fmt"
+
+	"github.com/rudderlabs/rudder-server/jobsdb"
 )
 
 type InvalidManager struct{}
 
+func (f *InvalidManager) Transform(job *jobsdb.JobT) (string, error) {
+	return "", errors.New("invalid job")
+}
+
 func (f *InvalidManager) Upload(asyncDestStruct *AsyncDestinationStruct) AsyncUploadOutput {
 	abortedJobIDs := append(asyncDestStruct.ImportingJobIDs, asyncDestStruct.FailedJobIDs...)
 	return AsyncUploadOutput{

diff --git a/router/batchrouter/asyncdestinationmanager/eloqua/bulk_uploader.go b/router/batchrouter/asyncdestinationmanager/eloqua/bulk_uploader.go
@@ -10,6 +10,7 @@ import (
 	"github.com/samber/lo"
 
 	"github.com/rudderlabs/rudder-go-kit/stats"
+	"github.com/rudderlabs/rudder-server/jobsdb"
 	"github.com/rudderlabs/rudder-server/router/batchrouter/asyncdestinationmanager/common"
 )
 
@@ -23,6 +24,10 @@ func (b *EloquaBulkUploader) createAsyncUploadErrorOutput(errorString string, er
 	}
 }
 
+func (b *EloquaBulkUploader) Transform(job *jobsdb.JobT) (string, error) {
+	return common.GetMarshalledData(job), nil
+}
+
 func (b *EloquaBulkUploader) Upload(asyncDestStruct *common.AsyncDestinationStruct) common.AsyncUploadOutput {
 	destination := asyncDestStruct.Destination
 	uploadRetryableStat := stats.Default.NewTaggedStat("events_over_prescribed_limit", stats.CountType, map[string]string{

diff --git a/router/batchrouter/asyncdestinationmanager/manager.go b/router/batchrouter/asyncdestinationmanager/manager.go
@@ -14,21 +14,6 @@ import (
 
 var json = jsoniter.ConfigCompatibleWithStandardLibrary
 
-func GetMarshalledData(payload string, jobID int64) string {
-	var job common.AsyncJob
-	err := json.Unmarshal([]byte(payload), &job.Message)
-	if err != nil {
-		panic("Unmarshalling Transformer Response Failed")
-	}
-	job.Metadata = make(map[string]interface{})
-	job.Metadata["job_id"] = jobID
-	responsePayload, err := json.Marshal(job)
-	if err != nil {
-		panic("Marshalling Response Payload Failed")
-	}
-	return string(responsePayload)
-}
-
 func NewManager(destination *backendconfig.DestinationT, backendConfig backendconfig.BackendConfig) (common.AsyncDestinationManager, error) {
 	switch destination.DestinationDefinition.Name {
 	case "BINGADS_AUDIENCE":

diff --git a/router/batchrouter/asyncdestinationmanager/marketo-bulk-upload/marketobulkupload.go b/router/batchrouter/asyncdestinationmanager/marketo-bulk-upload/marketobulkupload.go
@@ -227,6 +227,10 @@ func extractJobStats(keyMap map[string]interface{}, importingJobIDs []int64, sta
 	return succesfulJobIDs, failedJobIDsTrans
 }
 
+func (b *MarketoBulkUploader) Transform(job *jobsdb.JobT) (string, error) {
+	return common.GetMarshalledData(job), nil
+}
+
 func (b *MarketoBulkUploader) Upload(asyncDestStruct *common.AsyncDestinationStruct) common.AsyncUploadOutput {
 	destination := asyncDestStruct.Destination
 	destinationID := destination.ID

diff --git a/router/batchrouter/handle_async.go b/router/batchrouter/handle_async.go
@@ -17,7 +17,6 @@ import (
 
 	"github.com/rudderlabs/rudder-go-kit/stats"
 	"github.com/rudderlabs/rudder-server/jobsdb"
-	"github.com/rudderlabs/rudder-server/router/batchrouter/asyncdestinationmanager"
 	"github.com/rudderlabs/rudder-server/router/batchrouter/asyncdestinationmanager/common"
 	"github.com/rudderlabs/rudder-server/router/rterror"
 	"github.com/rudderlabs/rudder-server/router/types"
@@ -465,13 +464,21 @@ func (brt *Handle) sendJobsToStorage(batchJobs BatchedJobs) {
 	writeAtBytes := brt.asyncDestinationStruct[destinationID].Size
 	allowedSize := brt.maxPayloadSizeInBytes
 	for _, job := range batchJobs.Jobs {
-		transformedData := common.GetTransformedData(job.EventPayload)
 		if brt.asyncDestinationStruct[destinationID].Count < brt.maxEventsInABatch &&
 			!brt.asyncDestinationStruct[destinationID].UploadInProgress &&
 			brt.asyncDestinationStruct[destinationID].Size < allowedSize {
-			fileData := asyncdestinationmanager.GetMarshalledData(transformedData, job.JobID)
+			fileData, err := brt.asyncDestinationStruct[destinationID].Manager.Transform(job)
+			if err != nil {
+				failedAsyncJobs := BatchedJobs{
+					Jobs:       []*jobsdb.JobT{job},
+					Connection: batchJobs.Connection,
+					TimeWindow: batchJobs.TimeWindow,
+				}
+				brt.updateJobStatus(&failedAsyncJobs, false, err, false)
+				continue
+			}
 			brt.asyncDestinationStruct[destinationID].Size = brt.asyncDestinationStruct[destinationID].Size + len([]byte(fileData+"\n"))
-			_, err := file.WriteAt([]byte(fileData+"\n"), int64(writeAtBytes))
+			_, err = file.WriteAt([]byte(fileData+"\n"), int64(writeAtBytes))
 			if err != nil {
 				panic(fmt.Errorf("BRT: %s: file write failed : %s", brt.destType, err.Error()))
 			}