locust 2.31.4.dev27__py3-none-any.whl → 2.31.4.dev47__py3-none-any.whl

locust/_version.py

@@ -14,7 +14,7 @@ __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
 
-__version__ = "2.31.4.dev27"
+__version__ = "2.31.4.dev47"
 version = __version__
-__version_tuple__ = (2, 31, 4, "dev27")
+__version_tuple__ = (2, 31, 4, "dev47")
 version_tuple = __version_tuple__

locust/contrib/postgres.py

@@ -0,0 +1,45 @@
+from locust import User, events
+
+import time
+
+import psycopg
+
+
+class PostgresClient:
+    def __init__(self, conn_string):
+        self.connection = psycopg.connect(conn_string)
+
+    def execute_query(self, query):
+        start_time = time.time()
+        try:
+            self.connection.execute(query)
+            response_time = int((time.time() - start_time) * 1000)
+            events.request.fire(
+                request_type="postgres_success",
+                name="execute_query",
+                response_time=response_time,
+                response_length=0,
+            )
+        except Exception as e:
+            response_time = int((time.time() - start_time) * 1000)
+            events.request.fire(
+                request_type="postgres_failure",
+                name="execute_query",
+                response_time=response_time,
+                response_length=0,
+                exception=e,
+            )
+
+    def close(self):
+        self.connection.close()
+
+
+class PostgresUser(User):
+    abstract = True
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.client = PostgresClient(conn_string=self.conn_string)
+
+    def on_stop(self):
+        self.client.close()