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,38 @@
+from locust import TaskSet, between, task
+from locust.contrib.postgres import PostgresUser
+
+import os
+import random
+
+
+class UserTasks(TaskSet):
+    @task
+    def run_select_query(self):
+        self.client.execute_query(
+            self.user.conn_string,
+            "SELECT * FROM loadtesting.invoice WHERE amount > 500",
+        )
+
+    @task(3)
+    def run_update_query(self):
+        random_amount = random.randint(1, 12)
+        self.client.execute_query(
+            self.user.conn_string,
+            f"UPDATE loadtesting.invoice SET amount={random_amount} WHERE amount < 10",
+        )
+
+
+class PostgresLocust(PostgresUser):
+    tasks = [UserTasks]
+    min_wait = 0
+    max_wait = 3
+    wait_time = between(min_wait, max_wait)
+
+    # Use environment variables or default values
+    PGHOST = os.getenv("PGHOST", "localhost")
+    PGPORT = os.getenv("PGPORT", "5432")
+    PGDATABASE = os.getenv("PGDATABASE", "loadtesting_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,49 @@
+from locust import TaskSet, User, events
+
+import time
+
+import psycopg
+
+
+def create_conn(conn_string):
+    return psycopg.connect(conn_string)
+
+
+def execute_query(conn_string, query):
+    db_conn = create_conn(conn_string)
+    return db_conn.cursor().execute(query)
+
+
+class PostgresClient:
+    def __getattr__(self, name):
+        def request_handler(*args, **kwargs):
+            start_time = time.time()
+            try:
+                execute_query(*args, **kwargs)
+                response_time = int((time.time() - start_time) * 1000)
+                events.request.fire(
+                    request_type="postgres_success",
+                    name=name,
+                    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=name,
+                    response_time=response_time,
+                    response_length=0,
+                    exception=e,
+                )
+                print(f"error: {e}")
+
+        return request_handler
+
+
+class PostgresUser(User):
+    abstract = True
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.client = PostgresClient()