add utility methods to get and set bot id in context

forta-network · Dec 20, 2023 · 3a11883 · 3a11883
1 parent 5e17dc6
commit 3a11883
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/utils/apiutils/principal.go b/utils/apiutils/principal.go
@@ -7,6 +7,7 @@ type contextKey int
 const (
 	contextKeyPrincipalAddressValue contextKey = 1
 	contextKeyBatchValue            contextKey = 2
+	contextKeyBotID                 contextKey = 3
 )
 
 func GetAddress(ctx context.Context) string {
@@ -24,3 +25,11 @@ func GetBatch(ctx context.Context) string {
 func SetBatch(ctx context.Context, batch string) context.Context {
 	return context.WithValue(ctx, contextKeyBatchValue, batch)
 }
+
+func GetBotID(ctx context.Context) string {
+	return ctx.Value(contextKeyBotID).(string)
+}
+
+func SetBotID(ctx context.Context, addr string) context.Context {
+	return context.WithValue(ctx, contextKeyBotID, addr)
+}