Skip to content

Commit 9274d15

Browse files
elpranselprans
committed
Fix automatic cleanup of cursors and statements
Current coding for automatic collection and closure of cursors and statements is using a callback through a weakref, but stores this weakref _inside_ the tracked object, thus making the weakref part of a potential reference cycle. Per Python spec, the callback will only be invoked if the weakref is still alive, which will not be the case if a cursor or a statement were part of a reference cycle. This leads to a failure to close the cursor on database side and a subsequent DuplicateCursorError if the id got reused for another cursor. This is easily fixed by storing the weakref alongside the cursor and statement garbage lists.
1 parent 4dd74f1 commit 9274d15

2 files changed

Lines changed: 31 additions & 16 deletions

File tree

‎postgresql/driver/pq3.py‎

Lines changed: 5 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -618,10 +618,7 @@ def __init__(self, cursor_id, wref = weakref.ref, ID = ID):
618618
self._pq_cursor_id = self.database.typio.encode(cursor_id)
619619
# If the cursor's id was generated, it should be garbage collected.
620620
if cursor_id == ID(self):
621-
garbage = self.database.pq.garbage_cursors.append
622-
cid = self._pq_cursor_id
623-
# Callback for closing the cursor on remote end.
624-
self._del = wref(self, lambda x: garbage(cid))
621+
self.database.pq.register_cursor(self, self._pq_cursor_id)
625622
self._quoted_cursor_id = '"' + cursor_id.replace('"', '""') + '"'
626623
self._init()
627624

@@ -630,11 +627,8 @@ def __iter__(self):
630627

631628
def close(self):
632629
if self.closed is False:
633-
self.database.pq.garbage_cursors.append(self._pq_cursor_id)
630+
self.database.pq.trash_cursor(self._pq_cursor_id)
634631
self.closed = True
635-
# Don't need the weakref anymore.
636-
if hasattr(self, '_del'):
637-
del self._del
638632

639633
def _ins(self, *args):
640634
return xact.Instruction(*args, asynchook = self.database._receive_async)
@@ -1280,10 +1274,8 @@ def __init__(self,
12801274
self._pq_statement_id = database.typio._encode(self.statement_id)[0]
12811275

12821276
if not statement_id:
1283-
garbage = database.pq.garbage_statements.append
1284-
sid = self._pq_statement_id
1285-
# Callback for closing the statement on remote end.
1286-
self._del = wref(self, lambda x: garbage(sid))
1277+
# Register statement on a connection to close it automatically on db end
1278+
database.pq.register_statement(self, self._pq_statement_id)
12871279

12881280
def __repr__(self):
12891281
return '<{mod}.{name}[{ci}] {state}>'.format(
@@ -1424,11 +1416,8 @@ def sql_parameter_types(self):
14241416

14251417
def close(self):
14261418
if self.closed is False:
1427-
self.database.pq.garbage_statements.append(self._pq_statement_id)
1419+
self.database.pq.trash_statement(self._pq_statement_id)
14281420
self.closed = True
1429-
# Don't need the weakref anymore.
1430-
if hasattr(self, '_del'):
1431-
del self._del
14321421

14331422
def _init(self):
14341423
"""

‎postgresql/protocol/client3.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
Protocol version 3.0 client and tools.
66
"""
77
import os
8+
import weakref
89
from .buffer import pq_message_stream
910
from . import element3 as element
1011
from . import xact3 as xact
@@ -473,6 +474,28 @@ def complete(self):
473474
# only remove the transaction if it's *not* fatal
474475
self.xact = None
475476

477+
def register_cursor(self, cursor, pq_cursor_id):
478+
trash = self.trash_cursor
479+
self.cursors[pq_cursor_id] = weakref.ref(cursor, lambda ref: trash(pq_cursor_id))
480+
481+
def trash_cursor(self, pq_cursor_id):
482+
try:
483+
del self.cursors[pq_cursor_id]
484+
except KeyError:
485+
pass
486+
self.garbage_cursors.append(pq_cursor_id)
487+
488+
def register_statement(self, statement, pq_statement_id):
489+
trash = self.trash_statement
490+
self.statements[pq_statement_id] = weakref.ref(statement, lambda ref: trash(pq_statement_id))
491+
492+
def trash_statement(self, pq_statement_id):
493+
try:
494+
del self.statements[pq_statement_id]
495+
except KeyError:
496+
pass
497+
self.garbage_statements.append(pq_statement_id)
498+
476499
def __str__(self):
477500
if hasattr(self, 'ssl_negotiation'):
478501
if self.ssl_negotiation is True:
@@ -501,6 +524,9 @@ def __init__(self, socket_factory, startup, password = b'',):
501524
element.Startup(startup), password
502525
)
503526

527+
self.cursors = {}
528+
self.statements = {}
529+
504530
self.garbage_statements = []
505531
self.garbage_cursors = []
506532

0 commit comments

Comments
 (0)