From 73b289180ef4c8e2a445e5147425d08cf228ee9c Mon Sep 17 00:00:00 2001
From: Devanathan Sabapathy <devsaba@amazon.com>
Date: Fri, 5 Jan 2024 23:47:43 +0000
Subject: [PATCH] setting application name to track the usage of Workload
 Replicator

---
 common/util.py          |  4 +++-
 core/extract/extract.py | 23 ++++++++++++++++++++++-
 core/replay/replay.py   | 22 ++++++++++++++++++++++
 3 files changed, 47 insertions(+), 2 deletions(-)

diff --git a/common/util.py b/common/util.py
index 2f65fe693..007afcb03 100644
--- a/common/util.py
+++ b/common/util.py
@@ -23,10 +23,12 @@ def db_connect(
     database=None,
     odbc_driver=None,
     drop_return=False,
+    app_name=None
 ):
     if interface == "psql":
         conn = redshift_connector.connect(
-            user=username, password=password, host=host, port=port, database=database
+            user=username, password=password, host=host, port=port,
+            database=database, application_name=app_name
         )
 
         # if drop_return is set, monkey patch driver to not store result set in memory
diff --git a/core/extract/extract.py b/core/extract/extract.py
index d6782b9db..b15bf1eea 100644
--- a/core/extract/extract.py
+++ b/core/extract/extract.py
@@ -9,7 +9,7 @@
 import common.config as config_helper
 import common.log as log_helper
 from common import aws_service as aws_service_helper
-from common.util import cluster_dict
+from common.util import cluster_dict, db_connect
 import core.extract.extractor as extractor
 
 logger = logging.getLogger("WorkloadReplicatorLogger")
@@ -62,6 +62,27 @@ def main():
     if not e.load_driver():
         sys.exit("Failed to load driver")
 
+    # setting application name for tracking
+    application = "WorkloadReplicator-Extract"
+    
+    host = config.get("source_cluster_endpoint").split(".")[0]
+    port = int(config.get("source_cluster_endpoint").split(":")[-1].split("/")[0])
+    DbUser = config.get("master_username")
+    DbName = config.get("source_cluster_endpoint").split("/")[-1]
+    region = config.get("region")
+    endpoint = config.get('source_cluster_endpoint').split(":")[0]
+
+    response = aws_service_helper.redshift_get_cluster_credentials(
+        user=DbUser,
+        database_name=DbName,
+        cluster_id=host,
+        region=region)
+    db_connect(host=endpoint,
+               port=port,
+               database=DbName,
+               password=response['DbPassword'],
+               username=response['DbUser'], app_name=application)
+
     # Run extract job
     (
         extraction_name,
diff --git a/core/replay/replay.py b/core/replay/replay.py
index 978828a5f..f42debeb6 100644
--- a/core/replay/replay.py
+++ b/core/replay/replay.py
@@ -18,6 +18,7 @@
     is_serverless,
     CredentialsException,
     bucket_dict,
+    db_connect
 )
 import report_gen
 from unload_sys_table import UnloadSysTable
@@ -95,6 +96,27 @@ def main():
         logger.info("No logs to replay, nothing to do.")
         sys.exit()
 
+    # setting application name for tracking
+    application = "WorkloadReplicator-Replay"
+    
+    host = config.get("target_cluster_endpoint").split(".")[0]
+    port = int(config.get("target_cluster_endpoint").split(":")[-1].split("/")[0])
+    DbUser = config.get("master_username")
+    DbName = config.get("target_cluster_endpoint").split("/")[-1]
+    region = config.get("target_cluster_region")
+    endpoint = config.get('target_cluster_endpoint').split(":")[0]
+
+    response = aws_service_helper.redshift_get_cluster_credentials(
+        user=DbUser,
+        database_name=DbName,
+        cluster_id=host,
+        region=region)
+    db_connect(host=endpoint,
+               port=port,
+               database=DbName,
+               password=response['DbPassword'],
+               username=response['DbUser'], app_name=application)
+
     # Actual replay
     logger.debug("Starting replay")
     complete = False