From fdd497960a4aae259735e3ffdb67282f13aee7bc Mon Sep 17 00:00:00 2001
From: Kaushik Iska <iska.kaushik@gmail.com>
Date: Tue, 28 Nov 2023 15:45:59 -0500
Subject: [PATCH] Prepend PEERDB_DEPLOYMENT_UID only if its set (#735)

---
 flow/connectors/utils/env.go | 11 +++++++++++
 flow/shared/constants.go     | 16 ++++++++++++----
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/flow/connectors/utils/env.go b/flow/connectors/utils/env.go
index 2911e3d8ef..6d3065dae4 100644
--- a/flow/connectors/utils/env.go
+++ b/flow/connectors/utils/env.go
@@ -45,3 +45,14 @@ func GetEnvInt(name string, defaultValue int) int {
 
 	return i
 }
+
+// GetEnvString returns the value of the environment variable with the given name
+// or defaultValue if the environment variable is not set.
+func GetEnvString(name string, defaultValue string) string {
+	val, ok := GetEnv(name)
+	if !ok {
+		return defaultValue
+	}
+
+	return val
+}
diff --git a/flow/shared/constants.go b/flow/shared/constants.go
index 48a6f8ae2f..a10d529189 100644
--- a/flow/shared/constants.go
+++ b/flow/shared/constants.go
@@ -2,7 +2,8 @@ package shared
 
 import (
 	"fmt"
-	"os"
+
+	"github.com/PeerDB-io/peer-flow/connectors/utils"
 )
 
 const (
@@ -34,13 +35,20 @@ const (
 const FetchAndChannelSize = 256 * 1024
 
 func GetPeerFlowTaskQueueName(taskQueueID TaskQueueID) (string, error) {
-	deploymentUID := os.Getenv("PEERDB_DEPLOYMENT_UID")
 	switch taskQueueID {
 	case PeerFlowTaskQueueID:
-		return deploymentUID + "-" + peerFlowTaskQueue, nil
+		return prependUIDToTaskQueueName(peerFlowTaskQueue), nil
 	case SnapshotFlowTaskQueueID:
-		return deploymentUID + "-" + snapshotFlowTaskQueue, nil
+		return prependUIDToTaskQueueName(snapshotFlowTaskQueue), nil
 	default:
 		return "", fmt.Errorf("unknown task queue id %d", taskQueueID)
 	}
 }
+
+func prependUIDToTaskQueueName(taskQueueName string) string {
+	deploymentUID := utils.GetEnvString("PEERDB_DEPLOYMENT_UID", "")
+	if deploymentUID == "" {
+		return taskQueueName
+	}
+	return fmt.Sprintf("%s-%s", deploymentUID, taskQueueName)
+}