setting application name to track the usage of Workload Replicator (#137

)

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

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,

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