diff --git a/examples/postgres/README.md b/examples/postgres/README.md new file mode 100644 index 0000000000..304186efaa --- /dev/null +++ b/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 diff --git a/examples/postgres/locustfile.py b/examples/postgres/locustfile.py new file mode 100644 index 0000000000..fca4f9eee9 --- /dev/null +++ b/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}" diff --git a/locust/contrib/postgres.py b/locust/contrib/postgres.py new file mode 100644 index 0000000000..7cd1ead938 --- /dev/null +++ b/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()