Skip to content

Commit 7f9ce2a

Browse files
Nothing4YouNothing4You
committed
Add constants and tests related to query timeouts
1 parent ee88d0f commit 7f9ce2a

File tree

4 files changed

+175
-0
lines changed

4 files changed

+175
-0
lines changed

‎pymysql/constants/ER.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -470,5 +470,10 @@
470470
WRONG_STRING_LENGTH = 1468
471471
ERROR_LAST = 1468
472472

473+
# MariaDB only
474+
STATEMENT_TIMEOUT = 1969
475+
476+
QUERY_TIMEOUT = 3024
477+
473478
# https://github.com/PyMySQL/PyMySQL/issues/607
474479
CONSTRAINT_FAILED = 4025

‎pymysql/tests/base.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,14 @@ def mysql_server_is(self, conn, version_tuple):
4949
)
5050
return server_version_tuple >= version_tuple
5151

52+
def get_mysql_vendor(self, conn):
53+
server_version = conn.get_server_info()
54+
55+
if "MariaDB" in server_version:
56+
return "mariadb"
57+
58+
return "mysql"
59+
5260
_connections = None
5361

5462
@property

‎pymysql/tests/test_SSCursor.py

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,19 @@
11
import sys
22

3+
import pytest
4+
35
try:
46
from pymysql.tests import base
57
import pymysql.cursors
68
from pymysql.constants import CLIENT
9+
import pymysql.constants.ER
710
except Exception:
811
# For local testing from top-level directory, without installing
912
sys.path.append("../pymysql")
1013
from pymysql.tests import base
1114
import pymysql.cursors
1215
from pymysql.constants import CLIENT
16+
import pymysql.constants.ER
1317

1418

1519
class TestSSCursor(base.PyMySQLTestCase):
@@ -122,6 +126,94 @@ def test_SSCursor(self):
122126
cursor.execute("DROP TABLE IF EXISTS tz_data")
123127
cursor.close()
124128

129+
def test_execution_time_limit(self):
130+
# this method is similarly implemented in test_cursor
131+
132+
conn = self.connect()
133+
134+
# table creation and filling is SSCursor only as it's not provided by self.setUp()
135+
self.safe_create_table(
136+
conn,
137+
"test",
138+
"create table test (data varchar(10))",
139+
)
140+
with conn.cursor() as cur:
141+
cur.execute(
142+
"insert into test (data) values "
143+
"('row1'), ('row2'), ('row3'), ('row4'), ('row5')"
144+
)
145+
conn.commit()
146+
147+
db_type = self.get_mysql_vendor(conn)
148+
149+
with conn.cursor(pymysql.cursors.SSCursor) as cur:
150+
# MySQL MAX_EXECUTION_TIME takes ms
151+
# MariaDB max_statement_time takes seconds as int/float, introduced in 10.1
152+
153+
# this will sleep 0.01 seconds per row
154+
if db_type == "mysql":
155+
sql = (
156+
"SELECT /*+ MAX_EXECUTION_TIME(2000) */ data, sleep(0.01) FROM test"
157+
)
158+
else:
159+
sql = "SET STATEMENT max_statement_time=2 FOR SELECT data, sleep(0.01) FROM test"
160+
161+
cur.execute(sql)
162+
# unlike Cursor, SSCursor returns a list of tuples here
163+
self.assertEqual(
164+
cur.fetchall(),
165+
[
166+
("row1", 0),
167+
("row2", 0),
168+
("row3", 0),
169+
("row4", 0),
170+
("row5", 0),
171+
],
172+
)
173+
174+
if db_type == "mysql":
175+
sql = (
176+
"SELECT /*+ MAX_EXECUTION_TIME(2000) */ data, sleep(0.01) FROM test"
177+
)
178+
else:
179+
sql = "SET STATEMENT max_statement_time=2 FOR SELECT data, sleep(0.01) FROM test"
180+
cur.execute(sql)
181+
self.assertEqual(cur.fetchone(), ("row1", 0))
182+
183+
# this discards the previous unfinished query and raises an
184+
# incomplete unbuffered query warning
185+
with pytest.warns(UserWarning):
186+
cur.execute("SELECT 1")
187+
self.assertEqual(cur.fetchone(), (1,))
188+
189+
# SSCursor will not read the EOF packet until we try to read
190+
# another row. Skipping this will raise an incomplete unbuffered
191+
# query warning in the next cur.execute().
192+
self.assertEqual(cur.fetchone(), None)
193+
194+
if db_type == "mysql":
195+
sql = "SELECT /*+ MAX_EXECUTION_TIME(1) */ data, sleep(1) FROM test"
196+
else:
197+
sql = "SET STATEMENT max_statement_time=0.001 FOR SELECT data, sleep(1) FROM test"
198+
with pytest.raises(pymysql.err.OperationalError) as cm:
199+
# in an unbuffered cursor the OperationalError may not show up
200+
# until fetching the entire result
201+
cur.execute(sql)
202+
cur.fetchall()
203+
204+
if db_type == "mysql":
205+
# this constant was only introduced in MySQL 5.7, not sure
206+
# what was returned before, may have been ER_QUERY_INTERRUPTED
207+
self.assertEqual(cm.value.args[0], pymysql.constants.ER.QUERY_TIMEOUT)
208+
else:
209+
self.assertEqual(
210+
cm.value.args[0], pymysql.constants.ER.STATEMENT_TIMEOUT
211+
)
212+
213+
# connection should still be fine at this point
214+
cur.execute("SELECT 1")
215+
self.assertEqual(cur.fetchone(), (1,))
216+
125217

126218
__all__ = ["TestSSCursor"]
127219

‎pymysql/tests/test_cursor.py

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22

33
from pymysql.tests import base
44
import pymysql.cursors
5+
import pymysql.constants.ER
6+
7+
import pytest
58

69

710
class CursorTest(base.PyMySQLTestCase):
@@ -19,6 +22,7 @@ def setUp(self):
1922
"insert into test (data) values "
2023
"('row1'), ('row2'), ('row3'), ('row4'), ('row5')"
2124
)
25+
conn.commit()
2226
cursor.close()
2327
self.test_connection = pymysql.connect(**self.databases[0])
2428
self.addCleanup(self.test_connection.close)
@@ -129,3 +133,69 @@ def test_executemany(self):
129133
)
130134
finally:
131135
cursor.execute("DROP TABLE IF EXISTS percent_test")
136+
137+
def test_execution_time_limit(self):
138+
# this method is similarly implemented in test_SScursor
139+
140+
conn = self.test_connection
141+
db_type = self.get_mysql_vendor(conn)
142+
143+
with conn.cursor(pymysql.cursors.Cursor) as cur:
144+
# MySQL MAX_EXECUTION_TIME takes ms
145+
# MariaDB max_statement_time takes seconds as int/float, introduced in 10.1
146+
147+
# this will sleep 0.01 seconds per row
148+
if db_type == "mysql":
149+
sql = (
150+
"SELECT /*+ MAX_EXECUTION_TIME(2000) */ data, sleep(0.01) FROM test"
151+
)
152+
else:
153+
sql = "SET STATEMENT max_statement_time=2 FOR SELECT data, sleep(0.01) FROM test"
154+
155+
cur.execute(sql)
156+
# unlike SSCursor, Cursor returns a tuple of tuples here
157+
self.assertEqual(
158+
cur.fetchall(),
159+
(
160+
("row1", 0),
161+
("row2", 0),
162+
("row3", 0),
163+
("row4", 0),
164+
("row5", 0),
165+
),
166+
)
167+
168+
if db_type == "mysql":
169+
sql = (
170+
"SELECT /*+ MAX_EXECUTION_TIME(2000) */ data, sleep(0.01) FROM test"
171+
)
172+
else:
173+
sql = "SET STATEMENT max_statement_time=2 FOR SELECT data, sleep(0.01) FROM test"
174+
cur.execute(sql)
175+
self.assertEqual(cur.fetchone(), ("row1", 0))
176+
177+
# this discards the previous unfinished query
178+
cur.execute("SELECT 1")
179+
self.assertEqual(cur.fetchone(), (1,))
180+
181+
if db_type == "mysql":
182+
sql = "SELECT /*+ MAX_EXECUTION_TIME(1) */ data, sleep(1) FROM test"
183+
else:
184+
sql = "SET STATEMENT max_statement_time=0.001 FOR SELECT data, sleep(1) FROM test"
185+
with pytest.raises(pymysql.err.OperationalError) as cm:
186+
# in a buffered cursor this should reliably raise an
187+
# OperationalError
188+
cur.execute(sql)
189+
190+
if db_type == "mysql":
191+
# this constant was only introduced in MySQL 5.7, not sure
192+
# what was returned before, may have been ER_QUERY_INTERRUPTED
193+
self.assertEqual(cm.value.args[0], pymysql.constants.ER.QUERY_TIMEOUT)
194+
else:
195+
self.assertEqual(
196+
cm.value.args[0], pymysql.constants.ER.STATEMENT_TIMEOUT
197+
)
198+
199+
# connection should still be fine at this point
200+
cur.execute("SELECT 1")
201+
self.assertEqual(cur.fetchone(), (1,))

0 commit comments

Comments
 (0)