Merge pull request #2836 from guel-codes/postgres-test-example

Add PostgresUser to examples

examples/postgres/README.md

@@ -0,0 +1,21 @@
+# Overview
+
+Read the instruction below for your specific database
+
+## PostgreSQL
+
+### How to run the test
+
+- Prerequisites:
+
+  - `psycopg3` - https://www.psycopg.org/psycopg3/docs/basic/install.html
+
+- Set your environment variables for:
+
+  - PGHOST
+  - PGPORT
+  - PGDATABASE
+  - PGUSER
+  - PGPASSWORD
+
+- Run locust as usual, see https://docs.locust.io/en/stable/quickstart.html

examples/postgres/locustfile.py

@@ -0,0 +1,29 @@
+from locust import TaskSet, between, task
+from locust.contrib.postgres import PostgresUser
+
+import os
+import random
+
+
+class PostgresLocust(PostgresUser):
+    @task
+    def run_select_query(self):
+        self.client.execute_query(
+            "SELECT * FROM loadtesting.invoice WHERE amount > 500",
+        )
+
+    @task
+    def run_update_query(self):
+        random_amount = random.randint(1, 12)
+        self.client.execute_query(
+            f"UPDATE loadtesting.invoice SET amount={random_amount} WHERE amount < 10",
+        )
+
+    # Use environment variables or default values
+    PGHOST = os.getenv("PGHOST", "localhost")
+    PGPORT = os.getenv("PGPORT", "5432")
+    PGDATABASE = os.getenv("PGDATABASE", "test_db")
+    PGUSER = os.getenv("PGUSER", "postgres")
+    PGPASSWORD = os.getenv("PGPASSWORD", "postgres")
+
+    conn_string = f"postgresql://{PGUSER}:{PGPASSWORD}@{PGHOST}:{PGPORT}/{PGDATABASE}"

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()