Skip to content

Commit 7be2622

Browse files
author
James William Pye
committed
Alter chunks to return builtins.tuple by default.
Add some comments along the way.
1 parent f64c121 commit 7be2622

2 files changed

Lines changed: 82 additions & 62 deletions

File tree

‎postgresql/driver/pq3.py‎

Lines changed: 80 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -132,33 +132,30 @@ def _init(self):
132132
"""
133133
# The local initialization for the specific cursor.
134134

135-
def __init__(self, cursor_id):
135+
def __init__(self, cursor_id, wref = weakref.ref, ID = ID):
136136
self.cursor_id = cursor_id
137137
if self.statement is not None:
138138
self._output = self.statement._output
139139
self._output_io = self.statement._output_io
140140
self._output_formats = self.statement._output_formats or ()
141141
self._output_attmap = self.statement._output_attmap
142142

143-
if self.cursor_id == ID(self):
144-
addgarbage = self.database.pq.garbage_cursors.append
145-
typio = self.database.typio
143+
self._pq_cursor_id = self.database.typio.encode(cursor_id)
144+
# If the cursor's id was generated, it should be garbage collected.
145+
if cursor_id == ID(self):
146+
garbage = self.database.pq.garbage_cursors.append
147+
cid = self._pq_cursor_id
146148
# Callback for closing the cursor on remote end.
147-
self._del = weakref.ref(
148-
self, lambda x: addgarbage(typio.encode(cursor_id))
149-
)
150-
self._quoted_cursor_id = '"' + self.cursor_id.replace('"', '""') + '"'
151-
self._pq_cursor_id = self.database.typio.encode(self.cursor_id)
149+
self._del = wref(self, lambda x: garbage(cid))
150+
self._quoted_cursor_id = '"' + cursor_id.replace('"', '""') + '"'
152151
self._init()
153152

154153
def __iter__(self):
155154
return self
156155

157156
def close(self):
158157
if self.closed is False:
159-
self.database.pq.garbage_cursors.append(
160-
self.database.typio.encode(self.cursor_id)
161-
)
158+
self.database.pq.garbage_cursors.append(self._pq_cursor_id)
162159
self.closed = True
163160
# Don't need the weakref anymore.
164161
if hasattr(self, '_del'):
@@ -271,12 +268,11 @@ def _process_tuple_chunk_Row(self, x,
271268
proc = process_chunk,
272269
from_seq = Row.from_sequence,
273270
):
274-
return proc(self._output_io, x, self._raise_column_tuple_error)
275-
#attmap = self._output_attmap
276-
#return [
277-
# from_seq(attmap, y)
278-
# for y in proc(self._output_io, x, self._raise_column_tuple_error)
279-
#]
271+
attmap = self._output_attmap
272+
return [
273+
from_seq(attmap, y)
274+
for y in proc(self._output_io, x, self._raise_column_tuple_error)
275+
]
280276

281277
# Process the elemnt.Tuple messages in `x` for chunks()
282278
def _process_tuple_chunk(self, x, proc = process_chunk):
@@ -355,8 +351,12 @@ class Chunks(Output, pg_api.Chunks):
355351
pass
356352

357353
##
358-
# FetchAll is a Chunks cursor that gets all the information
359-
# in the cursor.
354+
# FetchAll - A Chunks cursor that gets *all* the records in the cursor.
355+
#
356+
# It has added complexity over other variants as in order to stream results,
357+
# chunks have to be removed from the protocol transaction's received messages.
358+
# If this wasn't done, the entire result set would be fully buffered prior
359+
# to processing.
360360
class FetchAll(Chunks):
361361
_e_factors = ('statement', 'parameters',)
362362
def _e_metas(self):
@@ -379,14 +379,17 @@ def _init(self,
379379
self._pq_xp_fetchall() + (element.SynchronizeMessage,)
380380
)
381381
self.database._pq_push(self._xact, self)
382+
383+
# Get more messages until the first Tuple is seen.
382384
STEP = self.database._pq_step
383385
while self._xact.state != xact.Complete:
384386
STEP()
385387
for x in self._xact.messages_received():
386388
if x.__class__ is tuple or expect == x.type:
387-
# no need to step once this is seen
389+
# No need to step anymore once this is seen.
388390
return
389391
elif x.type == null:
392+
# The protocol transaction is going to be complete..
390393
self.database._pq_complete()
391394
self._xact = None
392395
return
@@ -395,12 +398,15 @@ def _init(self,
395398
self.database._pq_complete()
396399
# If this was a select/copy cursor,
397400
# the data messages would have caused an earlier
398-
# return.
401+
# return. It's empty.
399402
self._xact = None
400403
return
401404
elif x.type in (bindcomplete, parsecomplete):
405+
# Noise.
402406
pass
403407
else:
408+
# This should have been caught by the protocol transaction.
409+
# "Can't happen".
404410
self.database._pq_complete()
405411
if self._xact.fatal is None:
406412
self._xact.fatal = False
@@ -412,15 +418,16 @@ def _init(self,
412418
self.database._raise_pq_error(self._xact, controller = self)
413419
return
414420

415-
def __next__(self):
421+
def __next__(self, data_types = (tuple,bytes)):
416422
x = self._xact
417423
# self._xact = None; means that the cursor has been exhausted.
418424
if x is None:
419425
raise StopIteration
420426

421427
# Finish the protocol transaction.
428+
STEP = self.database._pq_step
422429
while x.state is not xact.Complete and not x.completed:
423-
self.database._pq_step()
430+
STEP()
424431

425432
# fatal is None == no error
426433
# fatal is True == dead connection
@@ -432,10 +439,14 @@ def __next__(self):
432439
if not x.completed:
433440
# Transaction has been cleaned out of completed? iterator is done.
434441
self._xact = None
442+
self.close()
435443
raise StopIteration
436444

437445
# Get the chunk to be processed.
438-
chunk = x.completed[0][1]
446+
chunk = [
447+
y for y in x.completed[0][1]
448+
if y.__class__ in data_types
449+
]
439450
r = self._process_chunk(chunk)
440451
# Remove it, it's been processed.
441452
del x.completed[0]
@@ -447,21 +458,15 @@ class SingleXactCopy(FetchAll):
447458

448459
class SingleXactFetch(FetchAll):
449460
_expect = element.Tuple.type
450-
_process_chunk_ = FetchAll._process_tuple_chunk_Row
451-
452-
def _process_chunk(self, x, tuple_type = tuple):
453-
return self._process_chunk_((
454-
y for y in x if y.__class__ is tuple
455-
))
456461

457462
class MultiXactStream(Chunks):
458-
chunksize = 1024 * 3
463+
chunksize = 1024 * 4
459464
# only tuple streams
460-
_process_chunk = Output._process_tuple_chunk_Row
465+
_process_chunk = Output._process_tuple_chunk
461466

462467
def _e_metas(self):
463468
yield ('chunksize', self.chunksize)
464-
yield ('type', type(self).__name__)
469+
yield ('type', self.__class__.__name__)
465470

466471
def __init__(self, statement, parameters, cursor_id):
467472
self.statement = statement
@@ -505,6 +510,7 @@ def __next__(self, tuple_type = tuple):
505510
else:
506511
# it's done.
507512
self._xact = None
513+
self.close()
508514
if not chunk:
509515
# chunk is empty, it's done *right* now.
510516
raise StopIteration
@@ -740,38 +746,38 @@ def _e_metas(self):
740746
yield ('sql_column_types', ct)
741747

742748
def clone(self):
743-
ps = type(self)(self.database, None, self.string)
749+
ps = self.__class__(self.database, None, self.string)
744750
ps._init()
745751
ps._fini()
746752
return ps
747753

748-
def __init__(self, database, statement_id, string):
754+
def __init__(self,
755+
database, statement_id, string,
756+
wref = weakref.ref
757+
):
749758
self.database = database
750-
self.statement_id = statement_id or ID(self)
751759
self.string = string
760+
self.statement_id = statement_id or ID(self)
752761
self._xact = None
753-
self._pq_statement_id = None
754762
self.closed = None
763+
self._pq_statement_id = database.typio._encode(self.statement_id)[0]
755764

756765
if not statement_id:
757-
addgarbage = database.pq.garbage_statements.append
758-
typio = database.typio
759-
sid = self.statement_id
766+
garbage = database.pq.garbage_statements.append
767+
sid = self._pq_statement_id
760768
# Callback for closing the statement on remote end.
761-
self._del = weakref.ref(
762-
self, lambda x: addgarbage(typio.encode(sid))
763-
)
769+
self._del = wref(self, lambda x: garbage(sid))
764770

765771
def __repr__(self):
766772
return '<{mod}.{name}[{ci}] {state}>'.format(
767-
mod = type(self).__module__,
768-
name = type(self).__name__,
773+
mod = self.__class__.__module__,
774+
name = self.__class__.__name__,
769775
ci = self.database.connector._pq_iri,
770776
state = self.state,
771777
)
772778

773-
def _pq_parameters(self, parameters):
774-
return process_tuple(
779+
def _pq_parameters(self, parameters, proc = process_tuple):
780+
return proc(
775781
self._input_io, parameters,
776782
self._raise_parameter_tuple_error
777783
)
@@ -915,9 +921,6 @@ def _init(self):
915921
the return as there may be things that can be done while waiting
916922
for the return. Use the _fini() to complete.
917923
"""
918-
self._pq_statement_id = self.database.typio._encode(
919-
self.statement_id
920-
)[0]
921924
if self.string is not None:
922925
q = self.database.typio._encode(str(self.string))[0]
923926
cmd = [
@@ -993,23 +996,27 @@ def __call__(self, *parameters):
993996
raise TypeError("statement requires %d parameters, given %d" %(
994997
len(self._input), len(parameters)
995998
))
999+
##
9961000
# get em' all!
9971001
if self._output is None:
9981002
# might be a copy.
9991003
c = SingleXactCopy(self, parameters)
10001004
else:
10011005
c = SingleXactFetch(self, parameters)
1006+
c._process_chunk = c._process_tuple_chunk_Row
10021007

10031008
# iff output is None, it's not a tuple returning query.
10041009
# however, if it's a copy, detect that fact by SingleXactCopy's
10051010
# immediate return after finding the copy begin message(no complete).
1006-
if self._output is None and c.command() is not None:
1007-
return (c.command(), c.count())
1008-
else:
1009-
r = []
1010-
for x in c:
1011-
r.extend(x)
1012-
return r
1011+
if self._output is None:
1012+
cmd = c.command()
1013+
if cmd is not None:
1014+
return (cmd, c.count())
1015+
# Returns rows, accumulate in a list.
1016+
r = []
1017+
for x in c:
1018+
r.extend(x)
1019+
return r
10131020

10141021
def declare(self, *parameters):
10151022
if self.closed is None:
@@ -1022,7 +1029,10 @@ def declare(self, *parameters):
10221029
return Cursor(self, parameters, self.database, None)
10231030

10241031
def rows(self, *parameters, **kw):
1025-
return chain.from_iterable(self.chunks(*parameters, **kw))
1032+
chunks = self.chunks(*parameters, **kw)
1033+
if chunks._output_io:
1034+
chunks._process_chunk = chunks._process_tuple_chunk_Row
1035+
return chain.from_iterable(chunks)
10261036
__iter__ = rows
10271037

10281038
def chunks(self, *parameters):
@@ -1033,10 +1043,12 @@ def chunks(self, *parameters):
10331043
raise TypeError("statement requires %d parameters, given %d" %(
10341044
len(self._input), len(parameters)
10351045
))
1046+
10361047
if self._output is None:
1048+
# It's *probably* a COPY.
10371049
return SingleXactCopy(self, parameters)
10381050
if self.database.pq.state == b'I':
1039-
# Currently, *not* in a transaction block, so
1051+
# Currently, *not* in a Transaction block, so
10401052
# DECLARE the statement WITH HOLD in order to allow
10411053
# access across transactions.
10421054
if self.string is not None:
@@ -1047,6 +1059,7 @@ def chunks(self, *parameters):
10471059
# This happens when statement_from_id is used.
10481060
return SingleXactFetch(self, parameters)
10491061
else:
1062+
# Likely, the best possible case. It gets to use Execute messages.
10501063
return MultiXactInsideBlock(self, parameters, None)
10511064

10521065
def column(self, *parameters, **kw):
@@ -1056,12 +1069,15 @@ def column(self, *parameters, **kw):
10561069

10571070
def first(self, *parameters):
10581071
if self.closed is None:
1072+
# Not fully initialized; assume interrupted.
10591073
self._fini()
10601074
if self._input is not None:
1075+
# Use a regular TypeError.
10611076
if len(parameters) != len(self._input):
10621077
raise TypeError("statement requires %d parameters, given %d" %(
10631078
len(self._input), len(parameters)
10641079
))
1080+
10651081
# Parameters? Build em'.
10661082
db = self.database
10671083

@@ -1084,10 +1100,12 @@ def first(self, *parameters):
10841100
),
10851101
# Get all
10861102
element.Execute(b'', 0xFFFFFFFF),
1103+
element.ClosePortal(b''),
10871104
element.SynchronizeMessage
10881105
),
10891106
asynchook = db._receive_async
10901107
)
1108+
# Push and complete protocol transaction.
10911109
db._pq_push(x, self)
10921110
db._pq_complete()
10931111

@@ -1102,6 +1120,7 @@ def first(self, *parameters):
11021120
return None
11031121

11041122
if len(self._output_io) > 1:
1123+
# Multiple columns, return a Row.
11051124
return Row.from_sequence(
11061125
self._output_attmap,
11071126
process_tuple(
@@ -1110,6 +1129,7 @@ def first(self, *parameters):
11101129
),
11111130
)
11121131
else:
1132+
# Single column output.
11131133
if xt[0] is None:
11141134
return None
11151135
io = self._output_io[0] or self.database.typio.decode

‎postgresql/python/functools.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ def __call__(self, r):
3333
# C implementation of the tuple processors.
3434
from ..port.optimized import process_tuple, process_chunk
3535
except ImportError:
36-
def process_tuple(procs, tup, exception_handler, len = len):
36+
def process_tuple(procs, tup, exception_handler, len = len, tuple = tuple):
3737
"""
3838
Call each item in `procs` with the corresponding
3939
item in `tup` returning the result as `type`.
@@ -60,7 +60,7 @@ def process_tuple(procs, tup, exception_handler, len = len):
6060
# relying on __context__
6161
exception_handler(procs, tup, i)
6262
raise RuntimeError("process_tuple exception handler failed to raise")
63-
return r
63+
return tuple(r)
6464

6565
def process_chunk(procs, tupc, fail, process_tuple = process_tuple):
6666
return [process_tuple(procs, x, fail) for x in tupc]

0 commit comments

Comments
 (0)