Add PostgresUser to examples

examples/database/README.md

@@ -0,0 +1,31 @@
+# Overview
+
+Read the instruction below for your specific database
+
+## PostgreSQL
+### How to run the test
+- Prerequisites:
+    If you do not have them installed already
+    - `pip3 install psycopg3` - https://www.psycopg.org/psycopg3/docs/basic/install.html
+
+    - `pip3 install locust` - https://docs.locust.io/en/stable/installation.html
+
+- Set your environment variables for: 
+    - PGHOST
+    - PGPORT
+    - PGDATABASE
+    - PGUSER
+    - PGPASSWORD
+
+    so they can be picked up from `postgres_test.py`
+
+- Also, update the methods of the UserTasks class to run queries that will hit your database.
+
+#### Run test from browser:
+- Open the terminal and run `locust -f postgres_test.py` this will give you a URL to follow. Likely `http://0.0.0.0:8089` from here, you can put in custom parameters to run your test
+
+#### Run test from terminal:
+
+- Insert this into your terminal `locust -f name_of_test_file.py --users=500 --spawn-rate=10 --run-time='15s' --autostart --autoquit 3`
+- For a full list of all the configuration options that can be used run `locust --help` and visit https://docs.locust.io/en/stable/configuration.html#configuration
+

examples/database/postgresql/base_postgres_user.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
+    client_class = PostgresClient
+
+    def __init__(self, *args, **kwargs):
+        self.client = self.client_class()

examples/database/postgresql/postgres_test.py

@@ -0,0 +1,39 @@
+from locust import TaskSet, between, task
+
+import os
+import random
+
+from base_postgres_user import PostgresUser
+
+
+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}"