|
1 | 1 | import sys
|
2 | 2 |
|
| 3 | +import pytest |
| 4 | + |
3 | 5 | try:
|
4 | 6 | from pymysql.tests import base
|
5 | 7 | import pymysql.cursors
|
6 | 8 | from pymysql.constants import CLIENT
|
| 9 | + import pymysql.constants.ER |
7 | 10 | except Exception:
|
8 | 11 | # For local testing from top-level directory, without installing
|
9 | 12 | sys.path.append("../pymysql")
|
10 | 13 | from pymysql.tests import base
|
11 | 14 | import pymysql.cursors
|
12 | 15 | from pymysql.constants import CLIENT
|
| 16 | + import pymysql.constants.ER |
13 | 17 |
|
14 | 18 |
|
15 | 19 | class TestSSCursor(base.PyMySQLTestCase):
|
@@ -122,6 +126,94 @@ def test_SSCursor(self):
|
122 | 126 | cursor.execute("DROP TABLE IF EXISTS tz_data")
|
123 | 127 | cursor.close()
|
124 | 128 |
|
| 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 | + |
125 | 217 |
|
126 | 218 | __all__ = ["TestSSCursor"]
|
127 | 219 |
|
|
0 commit comments