From cf958a6ef6b2b36291561df19294255f9f884adc Mon Sep 17 00:00:00 2001 From: Miguel Johnson Date: Sun, 21 Jul 2024 18:46:47 -0400 Subject: [PATCH 01/14] Adding a PostgreSQL load test example --- examples/database/README.md | 21 +++++++++ examples/database/postgres_test.py | 73 ++++++++++++++++++++++++++++++ 2 files changed, 94 insertions(+) create mode 100644 examples/database/README.md create mode 100644 examples/database/postgres_test.py diff --git a/examples/database/README.md b/examples/database/README.md new file mode 100644 index 0000000000..3f7022c55b --- /dev/null +++ b/examples/database/README.md @@ -0,0 +1,21 @@ +# 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 psycopg2` - https://pypi.org/project/psycopg2/ + + - `pip3 install locust` - https://docs.locust.io/en/stable/installation.html + +- Update line 47 of the `postgres_test.py` file, you can change `postgresql://postgres:postgres@localhost:5432/loadtesting_db` to the connection string for 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 postgres_test.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 + diff --git a/examples/database/postgres_test.py b/examples/database/postgres_test.py new file mode 100644 index 0000000000..e5d34617e8 --- /dev/null +++ b/examples/database/postgres_test.py @@ -0,0 +1,73 @@ +import psycopg2 +from locust import User, TaskSet, task, between, events +import random, time + + +def create_conn(conn_string): + print("Connect and Query PostgreSQL") + conn = psycopg2.connect(conn_string) + return conn + + +def execute_query(conn_string, query): + db_conn = create_conn(conn_string) + db_query = db_conn.cursor().execute(query) + return db_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("error {}".format(e)) + + return request_handler + + +class UserTasks(TaskSet): + conn_string = "postgresql://postgres:postgres@localhost:5432/loadtesting_db" + + @task + def run_select_query(self): + self.client.execute_query( + self.conn_string, + f"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.conn_string, + f"UPDATE loadtesting.invoice SET amount={random_amount} WHERE amount < 10", + ) + + + +class PostgresLocust(User): + tasks = [UserTasks] + min_wait = 0 + max_wait = 3 + wait_time = between(min_wait, max_wait) + + def __init__(self, *args): + self.client = PostgresClient() From 08ef825627d7231dd8bcb344671ce562cc0799ad Mon Sep 17 00:00:00 2001 From: Miguel Johnson Date: Mon, 5 Aug 2024 14:10:50 -0400 Subject: [PATCH 02/14] updating to psycopg3 --- examples/database/README.md | 2 +- examples/database/postgres_test.py | 9 ++++++--- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/examples/database/README.md b/examples/database/README.md index 3f7022c55b..53112f6df0 100644 --- a/examples/database/README.md +++ b/examples/database/README.md @@ -6,7 +6,7 @@ Read the instruction below for your specific database ### How to run the test - Prerequisites: If you do not have them installed already - - `pip3 install psycopg2` - https://pypi.org/project/psycopg2/ + - `pip3 install psycopg3` - https://www.psycopg.org/psycopg3/docs/basic/install.html - `pip3 install locust` - https://docs.locust.io/en/stable/installation.html diff --git a/examples/database/postgres_test.py b/examples/database/postgres_test.py index e5d34617e8..58ff41bb47 100644 --- a/examples/database/postgres_test.py +++ b/examples/database/postgres_test.py @@ -1,6 +1,9 @@ -import psycopg2 -from locust import User, TaskSet, task, between, events -import random, time +import random +import time + +import psycopg + +from locust import TaskSet, User, between, events, task def create_conn(conn_string): From 548eb603caa8cb5e5aae2ef0bfa913c3d376e0ea Mon Sep 17 00:00:00 2001 From: Miguel Johnson Date: Mon, 5 Aug 2024 14:19:29 -0400 Subject: [PATCH 03/14] updating to psycopg3 --- examples/database/README.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/examples/database/README.md b/examples/database/README.md index 53112f6df0..ffe8f99173 100644 --- a/examples/database/README.md +++ b/examples/database/README.md @@ -10,7 +10,10 @@ Read the instruction below for your specific database - `pip3 install locust` - https://docs.locust.io/en/stable/installation.html -- Update line 47 of the `postgres_test.py` file, you can change `postgresql://postgres:postgres@localhost:5432/loadtesting_db` to the connection string for your database +- Update `line 50` in the UserTasks class of the `postgres_test.py` file, you can change `postgresql://postgres:postgres@localhost:5432/loadtesting_db` to the connection string for your database + +- 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 From d919018f95fd0bc50b56a33ecef0fbb3317d77cb Mon Sep 17 00:00:00 2001 From: Miguel Johnson Date: Mon, 5 Aug 2024 14:21:26 -0400 Subject: [PATCH 04/14] updating to psycopg3 --- examples/database/postgres_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/database/postgres_test.py b/examples/database/postgres_test.py index 58ff41bb47..d43f624778 100644 --- a/examples/database/postgres_test.py +++ b/examples/database/postgres_test.py @@ -8,7 +8,7 @@ def create_conn(conn_string): print("Connect and Query PostgreSQL") - conn = psycopg2.connect(conn_string) + conn = psycopg.connect(conn_string) return conn From 86aa6efd2bafeb75f62952c0c8e995bda5ba32a9 Mon Sep 17 00:00:00 2001 From: Miguel Johnson Date: Mon, 5 Aug 2024 14:33:25 -0400 Subject: [PATCH 05/14] fixing syntax errors for failed tests --- examples/database/postgres_test.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/database/postgres_test.py b/examples/database/postgres_test.py index d43f624778..76f7d1a73d 100644 --- a/examples/database/postgres_test.py +++ b/examples/database/postgres_test.py @@ -2,8 +2,8 @@ import time import psycopg +from locust import between, events, task, TaskSet, User -from locust import TaskSet, User, between, events, task def create_conn(conn_string): @@ -41,7 +41,7 @@ def request_handler(*args, **kwargs): exception=e, ) - print("error {}".format(e)) + print(f"error: {e}") return request_handler @@ -49,11 +49,11 @@ def request_handler(*args, **kwargs): class UserTasks(TaskSet): conn_string = "postgresql://postgres:postgres@localhost:5432/loadtesting_db" - @task + @task def run_select_query(self): self.client.execute_query( self.conn_string, - f"SELECT * FROM loadtesting.invoice WHERE amount > 500", + "SELECT * FROM loadtesting.invoice WHERE amount > 500", ) @task(3) From 01674a1a409933d670517a40510c7b70fd800adc Mon Sep 17 00:00:00 2001 From: Miguel Johnson Date: Mon, 5 Aug 2024 14:36:31 -0400 Subject: [PATCH 06/14] fixing syntax errors for failed tests --- examples/database/postgres_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/database/postgres_test.py b/examples/database/postgres_test.py index 76f7d1a73d..8a08f081cd 100644 --- a/examples/database/postgres_test.py +++ b/examples/database/postgres_test.py @@ -2,8 +2,8 @@ import time import psycopg -from locust import between, events, task, TaskSet, User +from locust import TaskSet, User, between, events, task def create_conn(conn_string): From b6737b7e8c717d4f05093be200cdf78b4bc151a9 Mon Sep 17 00:00:00 2001 From: Miguel Johnson Date: Mon, 5 Aug 2024 14:40:29 -0400 Subject: [PATCH 07/14] fixing syntax errors for failed tests --- examples/database/postgres_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/database/postgres_test.py b/examples/database/postgres_test.py index 8a08f081cd..b8d07f93ce 100644 --- a/examples/database/postgres_test.py +++ b/examples/database/postgres_test.py @@ -3,7 +3,7 @@ import psycopg -from locust import TaskSet, User, between, events, task +from locust import between, events, task, TaskSet, User def create_conn(conn_string): From 8e9663ffb038e43baed092dd13610187a39bbc25 Mon Sep 17 00:00:00 2001 From: Miguel Johnson Date: Mon, 5 Aug 2024 14:46:49 -0400 Subject: [PATCH 08/14] fixing imports --- examples/database/postgres_test.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/database/postgres_test.py b/examples/database/postgres_test.py index b8d07f93ce..ca319cfeea 100644 --- a/examples/database/postgres_test.py +++ b/examples/database/postgres_test.py @@ -1,10 +1,10 @@ +from locust import TaskSet, User, between, events, task + import random import time import psycopg -from locust import between, events, task, TaskSet, User - def create_conn(conn_string): print("Connect and Query PostgreSQL") From cd04c1509710c50bc12f399d44ee796037778f95 Mon Sep 17 00:00:00 2001 From: Miguel Johnson Date: Mon, 5 Aug 2024 14:50:49 -0400 Subject: [PATCH 09/14] small format fix --- examples/database/postgres_test.py | 1 - 1 file changed, 1 deletion(-) diff --git a/examples/database/postgres_test.py b/examples/database/postgres_test.py index ca319cfeea..59e010c445 100644 --- a/examples/database/postgres_test.py +++ b/examples/database/postgres_test.py @@ -65,7 +65,6 @@ def run_update_query(self): ) - class PostgresLocust(User): tasks = [UserTasks] min_wait = 0 From 516ba58bc14200e12704273b2471b397811cb3c6 Mon Sep 17 00:00:00 2001 From: Miguel Johnson Date: Mon, 19 Aug 2024 23:34:53 -0400 Subject: [PATCH 10/14] split out postgres test files --- examples/database/README.md | 11 ++- examples/database/postgres_test.py | 75 ------------------- .../database/postgresql/base_postgres_user.py | 49 ++++++++++++ examples/database/postgresql/postgres_test.py | 39 ++++++++++ 4 files changed, 97 insertions(+), 77 deletions(-) delete mode 100644 examples/database/postgres_test.py create mode 100644 examples/database/postgresql/base_postgres_user.py create mode 100644 examples/database/postgresql/postgres_test.py diff --git a/examples/database/README.md b/examples/database/README.md index ffe8f99173..a1b3121cfc 100644 --- a/examples/database/README.md +++ b/examples/database/README.md @@ -10,7 +10,14 @@ Read the instruction below for your specific database - `pip3 install locust` - https://docs.locust.io/en/stable/installation.html -- Update `line 50` in the UserTasks class of the `postgres_test.py` file, you can change `postgresql://postgres:postgres@localhost:5432/loadtesting_db` to the connection string for your database +- 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. @@ -19,6 +26,6 @@ Read the instruction below for your specific database #### Run test from terminal: -- Insert this into your terminal `locust -f postgres_test.py --users=500 --spawn-rate=10 --run-time='15s' --autostart --autoquit 3` +- 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 diff --git a/examples/database/postgres_test.py b/examples/database/postgres_test.py deleted file mode 100644 index 59e010c445..0000000000 --- a/examples/database/postgres_test.py +++ /dev/null @@ -1,75 +0,0 @@ -from locust import TaskSet, User, between, events, task - -import random -import time - -import psycopg - - -def create_conn(conn_string): - print("Connect and Query PostgreSQL") - conn = psycopg.connect(conn_string) - return conn - - -def execute_query(conn_string, query): - db_conn = create_conn(conn_string) - db_query = db_conn.cursor().execute(query) - return db_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 UserTasks(TaskSet): - conn_string = "postgresql://postgres:postgres@localhost:5432/loadtesting_db" - - @task - def run_select_query(self): - self.client.execute_query( - self.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.conn_string, - f"UPDATE loadtesting.invoice SET amount={random_amount} WHERE amount < 10", - ) - - -class PostgresLocust(User): - tasks = [UserTasks] - min_wait = 0 - max_wait = 3 - wait_time = between(min_wait, max_wait) - - def __init__(self, *args): - self.client = PostgresClient() diff --git a/examples/database/postgresql/base_postgres_user.py b/examples/database/postgresql/base_postgres_user.py new file mode 100644 index 0000000000..771a600734 --- /dev/null +++ b/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() diff --git a/examples/database/postgresql/postgres_test.py b/examples/database/postgresql/postgres_test.py new file mode 100644 index 0000000000..8186672347 --- /dev/null +++ b/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}" From 5d3b1ef4d5d940b7f83fa04340688d9ff2a7a6f8 Mon Sep 17 00:00:00 2001 From: Miguel Johnson Date: Tue, 20 Aug 2024 09:22:02 -0400 Subject: [PATCH 11/14] Addressing PR comments --- examples/database/README.md | 30 +++++++------------ examples/database/postgresql/postgres_test.py | 3 +- .../contrib/postgres.py | 4 +-- 3 files changed, 13 insertions(+), 24 deletions(-) rename examples/database/postgresql/base_postgres_user.py => locust/contrib/postgres.py (94%) diff --git a/examples/database/README.md b/examples/database/README.md index a1b3121cfc..304186efaa 100644 --- a/examples/database/README.md +++ b/examples/database/README.md @@ -3,29 +3,19 @@ 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` +### How to run the test -- Also, update the methods of the UserTasks class to run queries that will hit your database. +- Prerequisites: -#### 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 + - `psycopg3` - https://www.psycopg.org/psycopg3/docs/basic/install.html -#### Run test from terminal: +- Set your environment variables for: -- 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 + - PGHOST + - PGPORT + - PGDATABASE + - PGUSER + - PGPASSWORD +- Run locust as usual, see https://docs.locust.io/en/stable/quickstart.html diff --git a/examples/database/postgresql/postgres_test.py b/examples/database/postgresql/postgres_test.py index 8186672347..ee7c6c10c7 100644 --- a/examples/database/postgresql/postgres_test.py +++ b/examples/database/postgresql/postgres_test.py @@ -1,10 +1,9 @@ from locust import TaskSet, between, task +from locust.contrib.postgres import PostgresUser import os import random -from base_postgres_user import PostgresUser - class UserTasks(TaskSet): @task diff --git a/examples/database/postgresql/base_postgres_user.py b/locust/contrib/postgres.py similarity index 94% rename from examples/database/postgresql/base_postgres_user.py rename to locust/contrib/postgres.py index 771a600734..88febedfcb 100644 --- a/examples/database/postgresql/base_postgres_user.py +++ b/locust/contrib/postgres.py @@ -43,7 +43,7 @@ def request_handler(*args, **kwargs): class PostgresUser(User): abstract = True - client_class = PostgresClient def __init__(self, *args, **kwargs): - self.client = self.client_class() + super().__init__(*args, **kwargs) + self.client = PostgresClient() From 83f70e4f4cb099d58ac676f968133fa95fd89b78 Mon Sep 17 00:00:00 2001 From: Miguel Johnson Date: Tue, 20 Aug 2024 11:38:55 -0400 Subject: [PATCH 12/14] moving locustfile to examples/postgres/ --- examples/{database => postgres}/README.md | 0 .../postgresql/postgres_test.py => postgres/locustfile.py} | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename examples/{database => postgres}/README.md (100%) rename examples/{database/postgresql/postgres_test.py => postgres/locustfile.py} (100%) diff --git a/examples/database/README.md b/examples/postgres/README.md similarity index 100% rename from examples/database/README.md rename to examples/postgres/README.md diff --git a/examples/database/postgresql/postgres_test.py b/examples/postgres/locustfile.py similarity index 100% rename from examples/database/postgresql/postgres_test.py rename to examples/postgres/locustfile.py From b43a0c7a7a169bd7f372d9f558227a5aede6d8f6 Mon Sep 17 00:00:00 2001 From: Miguel Johnson Date: Tue, 20 Aug 2024 21:58:26 -0400 Subject: [PATCH 13/14] change how the connection is opened --- examples/postgres/locustfile.py | 4 +-- locust/contrib/postgres.py | 63 +++++++++++++++++---------------- 2 files changed, 34 insertions(+), 33 deletions(-) diff --git a/examples/postgres/locustfile.py b/examples/postgres/locustfile.py index ee7c6c10c7..6904ec88b0 100644 --- a/examples/postgres/locustfile.py +++ b/examples/postgres/locustfile.py @@ -9,7 +9,6 @@ class UserTasks(TaskSet): @task def run_select_query(self): self.client.execute_query( - self.user.conn_string, "SELECT * FROM loadtesting.invoice WHERE amount > 500", ) @@ -17,7 +16,6 @@ def run_select_query(self): 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", ) @@ -31,7 +29,7 @@ class PostgresLocust(PostgresUser): # Use environment variables or default values PGHOST = os.getenv("PGHOST", "localhost") PGPORT = os.getenv("PGPORT", "5432") - PGDATABASE = os.getenv("PGDATABASE", "loadtesting_db") + PGDATABASE = os.getenv("PGDATABASE", "postgres") PGUSER = os.getenv("PGUSER", "postgres") PGPASSWORD = os.getenv("PGPASSWORD", "postgres") diff --git a/locust/contrib/postgres.py b/locust/contrib/postgres.py index 88febedfcb..7444666dd0 100644 --- a/locust/contrib/postgres.py +++ b/locust/contrib/postgres.py @@ -9,36 +9,36 @@ 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 + def __init__(self, conn_string): + self.conn_string = conn_string + self.connection = create_conn(conn_string) + + def execute_query(self, query): + start_time = time.time() + try: + cursor = self.connection.cursor() + cursor.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, + ) + print(f"error: {e}") + + def close(self): + self.connection.close() class PostgresUser(User): @@ -46,4 +46,7 @@ class PostgresUser(User): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.client = PostgresClient() + self.client = PostgresClient(conn_string=self.conn_string) + + def on_stop(self): + self.client.close() From 4931df14b7528a79a6e393358b0fc6c1eb71afea Mon Sep 17 00:00:00 2001 From: Miguel Johnson Date: Sat, 24 Aug 2024 09:01:44 -0400 Subject: [PATCH 14/14] clean up locustfile --- examples/postgres/locustfile.py | 13 +++---------- locust/contrib/postgres.py | 13 +++---------- 2 files changed, 6 insertions(+), 20 deletions(-) diff --git a/examples/postgres/locustfile.py b/examples/postgres/locustfile.py index 6904ec88b0..fca4f9eee9 100644 --- a/examples/postgres/locustfile.py +++ b/examples/postgres/locustfile.py @@ -5,31 +5,24 @@ import random -class UserTasks(TaskSet): +class PostgresLocust(PostgresUser): @task def run_select_query(self): self.client.execute_query( "SELECT * FROM loadtesting.invoice WHERE amount > 500", ) - @task(3) + @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", ) - -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", "postgres") + PGDATABASE = os.getenv("PGDATABASE", "test_db") PGUSER = os.getenv("PGUSER", "postgres") PGPASSWORD = os.getenv("PGPASSWORD", "postgres") diff --git a/locust/contrib/postgres.py b/locust/contrib/postgres.py index 7444666dd0..7cd1ead938 100644 --- a/locust/contrib/postgres.py +++ b/locust/contrib/postgres.py @@ -1,24 +1,18 @@ -from locust import TaskSet, User, events +from locust import User, events import time import psycopg -def create_conn(conn_string): - return psycopg.connect(conn_string) - - class PostgresClient: def __init__(self, conn_string): - self.conn_string = conn_string - self.connection = create_conn(conn_string) + self.connection = psycopg.connect(conn_string) def execute_query(self, query): start_time = time.time() try: - cursor = self.connection.cursor() - cursor.execute(query) + self.connection.execute(query) response_time = int((time.time() - start_time) * 1000) events.request.fire( request_type="postgres_success", @@ -35,7 +29,6 @@ def execute_query(self, query): response_length=0, exception=e, ) - print(f"error: {e}") def close(self): self.connection.close()