From 9b61411c39c177ee1525678f0888f4ba706d746d Mon Sep 17 00:00:00 2001
From: Lars Holmberg <lars.holmberg@svenskaspel.se>
Date: Tue, 2 Nov 2021 16:04:18 +0100
Subject: [PATCH] Timescale listener: Save url from request event to database

---
 locust_plugins/listeners.py         | 16 +++++++++++++---
 locust_plugins/timescale_schema.sql |  3 ++-
 2 files changed, 15 insertions(+), 4 deletions(-)

diff --git a/locust_plugins/listeners.py b/locust_plugins/listeners.py
index 4938163..8cb6878 100644
--- a/locust_plugins/listeners.py
+++ b/locust_plugins/listeners.py
@@ -150,9 +150,9 @@ def write_samples_to_db(self, samples):
             with self._conn.cursor() as cur:
                 psycopg2.extras.execute_values(
                     cur,
-                    """INSERT INTO request(time,run_id,greenlet_id,loadgen,name,request_type,response_time,success,testplan,response_length,exception,pid,context) VALUES %s""",
+                    """INSERT INTO request(time,run_id,greenlet_id,loadgen,name,request_type,response_time,success,testplan,response_length,exception,pid,url,context) VALUES %s""",
                     samples,
-                    template="(%(time)s, %(run_id)s, %(greenlet_id)s, %(loadgen)s, %(name)s, %(request_type)s, %(response_time)s, %(success)s, %(testplan)s, %(response_length)s, %(exception)s, %(pid)s, %(context)s)",
+                    template="(%(time)s, %(run_id)s, %(greenlet_id)s, %(loadgen)s, %(name)s, %(request_type)s, %(response_time)s, %(success)s, %(testplan)s, %(response_length)s, %(exception)s, %(pid)s, %(url)s, %(context)s)",
                 )
         except psycopg2.Error as error:
             logging.error("Failed to write samples to Postgresql timescale database: " + repr(error))
@@ -166,7 +166,16 @@ def quitting(self, **_kwargs):
         self.exit()
 
     def on_request(
-        self, request_type, name, response_time, response_length, exception, context, start_time=None, **_kwargs
+        self,
+        request_type,
+        name,
+        response_time,
+        response_length,
+        exception,
+        context,
+        start_time=None,
+        url=None,
+        **_kwargs,
     ):
         success = 0 if exception else 1
         if start_time:
@@ -185,6 +194,7 @@ def on_request(
             "request_type": request_type,
             "response_time": response_time,
             "success": success,
+            "url": url[0:255] if url else None,
             "testplan": self._testplan,
             "pid": self._pid,
             "context": psycopg2.extras.Json(context, safe_serialize),
diff --git a/locust_plugins/timescale_schema.sql b/locust_plugins/timescale_schema.sql
index 11dfe39..33a5904 100644
--- a/locust_plugins/timescale_schema.sql
+++ b/locust_plugins/timescale_schema.sql
@@ -62,7 +62,8 @@ CREATE TABLE public.request (
     success smallint NOT NULL,
     testplan character varying(30) NOT NULL,
     pid integer,
-    context jsonb
+    context jsonb,
+    url character varying(255)
 );