add table mapping for metrics in eventhub

PeerDB-io · Sep 12, 2023 · 7e2309c · 7e2309c
1 parent d767df0
commit 7e2309c
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/flow/connectors/eventhub/eventhub.go b/flow/connectors/eventhub/eventhub.go
@@ -112,7 +112,7 @@ func (c *EventHubConnector) SyncRecords(req *model.SyncRecordsRequest) (*model.S
 	defer func() {
 		shutdown <- true
 	}()
-
+	tableNameRowsMapping := make(map[string]uint32)
 	batch := req.Records
 	eventsPerHeartBeat := 1000
 	eventsPerBatch := int(req.PushBatchSize)
@@ -149,7 +149,7 @@ func (c *EventHubConnector) SyncRecords(req *model.SyncRecordsRequest) (*model.S
 
 		if (i+1)%eventsPerBatch == 0 {
 			err := c.sendEventBatch(batchPerTopic, maxParallelism,
-				req.FlowJobName)
+				req.FlowJobName, tableNameRowsMapping)
 			if err != nil {
 				return nil, err
 			}
@@ -161,7 +161,7 @@ func (c *EventHubConnector) SyncRecords(req *model.SyncRecordsRequest) (*model.S
 	// send the remaining events.
 	if len(batchPerTopic) > 0 {
 		err := c.sendEventBatch(batchPerTopic, maxParallelism,
-			req.FlowJobName)
+			req.FlowJobName, tableNameRowsMapping)
 		if err != nil {
 			return nil, err
 		}
@@ -186,7 +186,8 @@ func (c *EventHubConnector) SyncRecords(req *model.SyncRecordsRequest) (*model.S
 
 func (c *EventHubConnector) sendEventBatch(events map[string][]*eventhub.Event,
 	maxParallelism int64,
-	flowName string) error {
+	flowName string,
+	tableNameRowsMapping map[string]uint32) error {
 	if len(events) == 0 {
 		log.WithFields(log.Fields{
 			"flowName": flowName,
@@ -233,6 +234,7 @@ func (c *EventHubConnector) sendEventBatch(events map[string][]*eventhub.Event,
 				"flowName": flowName,
 			}).Infof("pushed %d events to event hub: %s",
 				numEventsPushed, tblName)
+			tableNameRowsMapping[tblName] += uint32(len(eventBatch))
 		}(tblName, eventBatch)
 	}