Skip to content

Commit

Permalink
Discard results without converting them into Python objects. (#601)
Browse files Browse the repository at this point in the history
Fixes #560.
  • Loading branch information
methane committed May 16, 2023
1 parent 0220f42 commit a2e9706
Show file tree
Hide file tree
Showing 3 changed files with 126 additions and 5 deletions.
69 changes: 69 additions & 0 deletions src/MySQLdb/_mysql.c
Original file line number Diff line number Diff line change
Expand Up @@ -1484,6 +1484,26 @@ _mysql_ResultObject_fetch_row(
return NULL;
}

static const char _mysql_ResultObject_discard__doc__[] =
"discard() -- Discard remaining rows in the resultset.";

static PyObject *
_mysql_ResultObject_discard(
_mysql_ResultObject *self,
PyObject *noargs)
{
check_result_connection(self);

MYSQL_ROW row;
while (NULL != (row = mysql_fetch_row(self->result))) {
// do nothing
}
if (mysql_errno(self->conn)) {
return _mysql_Exception(self->conn);
}
Py_RETURN_NONE;
}

static char _mysql_ConnectionObject_change_user__doc__[] =
"Changes the user and causes the database specified by db to\n\
become the default (current) database on the connection\n\
Expand Down Expand Up @@ -2081,6 +2101,43 @@ _mysql_ConnectionObject_use_result(
return result;
}

static const char _mysql_ConnectionObject_discard_result__doc__[] =
"Discard current result set.\n\n"
"This function can be called instead of use_result() or store_result(). Non-standard.";

static PyObject *
_mysql_ConnectionObject_discard_result(
_mysql_ConnectionObject *self,
PyObject *noargs)
{
check_connection(self);
MYSQL *conn = &(self->connection);

Py_BEGIN_ALLOW_THREADS;

MYSQL_RES *res = mysql_use_result(conn);
if (res == NULL) {
Py_BLOCK_THREADS;
if (mysql_errno(conn) != 0) {
// fprintf(stderr, "mysql_use_result failed: %s\n", mysql_error(conn));
return _mysql_Exception(self);
}
Py_RETURN_NONE;
}

MYSQL_ROW row;
while (NULL != (row = mysql_fetch_row(res))) {
// do nothing.
}
mysql_free_result(res);
Py_END_ALLOW_THREADS;
if (mysql_errno(conn)) {
// fprintf(stderr, "mysql_free_result failed: %s\n", mysql_error(conn));
return _mysql_Exception(self);
}
Py_RETURN_NONE;
}

static void
_mysql_ConnectionObject_dealloc(
_mysql_ConnectionObject *self)
Expand Down Expand Up @@ -2376,6 +2433,12 @@ static PyMethodDef _mysql_ConnectionObject_methods[] = {
METH_NOARGS,
_mysql_ConnectionObject_use_result__doc__
},
{
"discard_result",
(PyCFunction)_mysql_ConnectionObject_discard_result,
METH_NOARGS,
_mysql_ConnectionObject_discard_result__doc__
},
{NULL, NULL} /* sentinel */
};

Expand Down Expand Up @@ -2437,6 +2500,12 @@ static PyMethodDef _mysql_ResultObject_methods[] = {
METH_VARARGS | METH_KEYWORDS,
_mysql_ResultObject_fetch_row__doc__
},
{
"discard",
(PyCFunction)_mysql_ResultObject_discard,
METH_NOARGS,
_mysql_ResultObject_discard__doc__
},
{
"field_flags",
(PyCFunction)_mysql_ResultObject_field_flags,
Expand Down
24 changes: 20 additions & 4 deletions src/MySQLdb/cursors.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,13 +75,30 @@ def __init__(self, connection):
self.rownumber = None
self._rows = None

def _discard(self):
self.description = None
self.description_flags = None
self.rowcount = -1
self.lastrowid = None
self._rows = None
self.rownumber = None

if self._result:
self._result.discard()
self._result = None

con = self.connection
if con is None:
return
while con.next_result() == 0: # -1 means no more data.
con.discard_result()

def close(self):
"""Close the cursor. No further queries will be possible."""
try:
if self.connection is None:
return
while self.nextset():
pass
self._discard()
finally:
self.connection = None
self._result = None
Expand Down Expand Up @@ -180,8 +197,7 @@ def execute(self, query, args=None):
Returns integer represents rows affected, if any
"""
while self.nextset():
pass
self._discard()

mogrified_query = self._mogrify(query, args)

Expand Down
38 changes: 37 additions & 1 deletion tests/test_cursor.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
# import pytest
import pytest
import MySQLdb.cursors
from configdb import connection_factory

Expand Down Expand Up @@ -186,3 +186,39 @@ def test_mogrify_with_dict_args():

assert mogrified_query == "SELECT 1, 2"
assert mogrified_query == cursor._executed.decode()


# Test that cursor can be used without reading whole resultset.
@pytest.mark.parametrize("Cursor", [MySQLdb.cursors.Cursor, MySQLdb.cursors.SSCursor])
def test_cursor_discard_result(Cursor):
conn = connect()
cursor = conn.cursor(Cursor)

cursor.execute(
"""\
CREATE TABLE test_cursor_discard_result (
id INTEGER PRIMARY KEY AUTO_INCREMENT,
data VARCHAR(100)
)"""
)
_tables.append("test_cursor_discard_result")

cursor.executemany(
"INSERT INTO test_cursor_discard_result (id, data) VALUES (%s, %s)",
[(i, f"row {i}") for i in range(1, 101)],
)

cursor.execute(
"""\
SELECT * FROM test_cursor_discard_result WHERE id <= 10;
SELECT * FROM test_cursor_discard_result WHERE id BETWEEN 11 AND 20;
SELECT * FROM test_cursor_discard_result WHERE id BETWEEN 21 AND 30;
"""
)
cursor.nextset()
assert cursor.fetchone() == (11, "row 11")

cursor.execute(
"SELECT * FROM test_cursor_discard_result WHERE id BETWEEN 31 AND 40"
)
assert cursor.fetchone() == (31, "row 31")

0 comments on commit a2e9706

Please sign in to comment.