posulliv · December 15, 2022 19:51 · Dec 15, 2022
diff --git a/locustfile.py b/locustfile.py
@@ -0,0 +1,82 @@
+import time
+
+from locust import User, TaskSet, task, between, events
+from trino.dbapi import connect
+from trino import auth
+from configparser import ConfigParser
+
+config = ConfigParser()
+config.read('config.ini')
+trino_config = config['trino']
+
+conn = connect(
+    host = trino_config['host'],
+    port = trino_config['port'],
+    user = trino_config['user'],
+    catalog = trino_config['catalog'],
+    schema = trino_config['schema'],
+    http_scheme = trino_config['http_scheme'],
+    auth = auth.BasicAuthentication(trino_config['user'], trino_config['password']),
+    verify = trino_config['verify_certs'].lower() == 'true'
+)
+
+query = """explain analyze SELECT sum(extendedprice * discount) AS revenue
+FROM
+  lineitem
+WHERE
+  shipdate >= DATE '1994-01-01'
+  AND shipdate < DATE '1994-01-01' + INTERVAL '1' YEAR
+  AND discount BETWEEN decimal '0.06' - decimal '0.01' AND decimal '0.06' + decimal '0.01'
+  AND quantity < 24"""
+
+def execute_query():
+
+    cur = conn.cursor()
+    cur.execute(query)
+    rows = cur.fetchall()
+    cur.close()
+
+    return rows
+
+class TrinoClient:
+
+    def __getattr__(self, name):
+        def wrapper(*args, **kwargs):
+            start_time = time.time()
+            try:
+                res = execute_query()
+                events.request_success.fire(
+                    request_type="trino",
+                    name=name,
+                    response_time=int((time.time() - start_time) * 1000),
+                    response_length=len(res)
+                )
+            except Exception as e:
+                events.request_failure.fire(
+                    request_type="trino",
+                    name=name,
+                    response_time=int((time.time() - start_time) * 1000),
+                    exception=e
+                )
+
+        return wrapper
+
+
+class TrinoTaskSet(TaskSet):
+    @task
+    def execute_query(self):
+        self.client.execute_query()
+
+
+class TrinoUser(User):
+    host = trino_config['host']
+
+    wait_time = between(0.1, 1)
+
+    def __init__(self, environment):
+        super(TrinoUser, self).__init__(environment)
+        self.client = TrinoClient()
+
+    @task
+    def execute_query(self):
+        self.client.execute_query()
diff --git a/requirements.txt b/requirements.txt
@@ -0,0 +1,2 @@
+locust==2.9.0
+trino==0.320.0
No results found