Skip to content

Commit 0a23467

Browse files
author
James William Pye
committed
Use a message_types tuple to avoid duplicate message types.
This allows for safe usage of the is operator in most message type comparisons. It allows reduces memory usage/allocation time(pbuffer still has to instantiate the integer for dereferencing the actual message_type, so there's probably no win there).
1 parent 3413eba commit 0a23467

5 files changed

Lines changed: 102 additions & 52 deletions

File tree

postgresql/protocol/buffer.c

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717
#include <Python.h>
1818
#include <structmember.h>
1919

20+
static PyObject *message_types = NULL;
21+
2022
struct p_list
2123
{
2224
PyObject *data; /* PyString pushed onto the buffer */
@@ -304,12 +306,16 @@ p_build_tuple(struct p_place *p)
304306
p_seek(p, copy_amount);
305307
}
306308

307-
mt = PyBytes_FromStringAndSize(header, 1);
309+
mt = PyTuple_GET_ITEM(message_types, (int) header[0]);
308310
if (mt == NULL)
309311
{
312+
/*
313+
* With message_types, this is nearly a can't happen.
314+
*/
310315
if (body != NULL) free(body);
311316
return(NULL);
312317
}
318+
Py_INCREF(mt);
313319

314320
md = PyBytes_FromStringAndSize(body, (Py_ssize_t) msg_length);
315321
if (body != NULL)
@@ -581,6 +587,8 @@ PyMODINIT_FUNC
581587
PyInit_cbuffer(void)
582588
{
583589
PyObject *mod;
590+
PyObject *msgtypes;
591+
PyObject *fromlist, *fromstr;
584592

585593
mod = PyModule_Create(&cbuffermodule);
586594
if (mod == NULL)
@@ -593,6 +601,31 @@ PyInit_cbuffer(void)
593601
(PyObject *) &pq_message_stream_Type) < 0)
594602
goto cleanup;
595603

604+
/*
605+
* Get the message_types tuple to type "instantiation".
606+
*/
607+
fromlist = PyList_New(1);
608+
fromstr = PyUnicode_FromString("message_types");
609+
PyList_SetItem(fromlist, 0, fromstr);
610+
msgtypes = PyImport_ImportModuleLevel(
611+
"message_types",
612+
PyModule_GetDict(mod),
613+
PyModule_GetDict(mod),
614+
fromlist, 1
615+
);
616+
Py_DECREF(fromlist);
617+
if (msgtypes == NULL)
618+
goto cleanup;
619+
message_types = PyObject_GetAttrString(msgtypes, "message_types");
620+
Py_DECREF(msgtypes);
621+
622+
if (!PyObject_IsInstance(message_types, (PyObject *) (&PyTuple_Type)))
623+
{
624+
PyErr_SetString(PyExc_RuntimeError,
625+
"local protocol.message_types.message_types is not a tuple object");
626+
goto cleanup;
627+
}
628+
596629
return(mod);
597630

598631
cleanup:

postgresql/protocol/element3.py

Lines changed: 42 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313
except ImportError:
1414
pass
1515

16+
from .message_types import message_types
17+
1618
StringFormat = b'\x00\x00'
1719
BinaryFormat = b'\x00\x01'
1820

@@ -115,8 +117,7 @@ def dict_message_repr(self):
115117

116118
class WireMessage(Message):
117119
def __init__(self, typ_data):
118-
type = bytes(type)[0]
119-
self.type = typ_data[0]
120+
self.type = message_types[typ_data[0][0]]
120121
self.data = typ_data[1]
121122

122123
def serialize(self):
@@ -131,7 +132,7 @@ def parse(typ, data):
131132
len(data), ulong.unpack(data[1:5])[0] + 1
132133
)
133134
)
134-
return typ((data[0], data[5:]))
135+
return typ((data[0:1], data[5:]))
135136

136137
class EmptyMessage(Message):
137138
'An abstract message that is always empty'
@@ -152,7 +153,7 @@ def parse(typ, data):
152153

153154
class Notify(Message):
154155
'Asynchronous notification message'
155-
type = b'A'
156+
type = message_types[b'A'[0]]
156157
__slots__ = ('pid', 'relation', 'parameter')
157158

158159
def __init__(self, pid, relation, parameter = b''):
@@ -174,7 +175,7 @@ def parse(typ, data):
174175
class ShowOption(Message):
175176
"""ShowOption(name, value)
176177
GUC variable information from backend"""
177-
type = b'S'
178+
type = message_types[b'S'[0]]
178179
__slots__ = ('name', 'value')
179180

180181
def __init__(self, name, value):
@@ -190,7 +191,7 @@ def parse(typ, data):
190191

191192
class Complete(StringMessage):
192193
'Command completion message.'
193-
type = b'C'
194+
type = message_types[b'C'[0]]
194195
__slots__ = ()
195196

196197
@classmethod
@@ -216,49 +217,49 @@ def extract_command(self):
216217

217218
class Null(EmptyMessage):
218219
'Null command'
219-
type = b'I'
220+
type = message_types[b'I'[0]]
220221
__slots__ = ()
221222
NullMessage = Message.__new__(Null)
222223
Null.SingleInstance = NullMessage
223224

224225
class NoData(EmptyMessage):
225226
'Null command'
226-
type = b'n'
227+
type = message_types[b'n'[0]]
227228
__slots__ = ()
228229
NoDataMessage = Message.__new__(NoData)
229230
NoData.SingleInstance = NoDataMessage
230231

231232
class ParseComplete(EmptyMessage):
232233
'Parse reaction'
233-
type = b'1'
234+
type = message_types[b'1'[0]]
234235
__slots__ = ()
235236
ParseCompleteMessage = Message.__new__(ParseComplete)
236237
ParseComplete.SingleInstance = ParseCompleteMessage
237238

238239
class BindComplete(EmptyMessage):
239240
'Bind reaction'
240-
type = b'2'
241+
type = message_types[b'2'[0]]
241242
__slots__ = ()
242243
BindCompleteMessage = Message.__new__(BindComplete)
243244
BindComplete.SingleInstance = BindCompleteMessage
244245

245246
class CloseComplete(EmptyMessage):
246247
'Close statement or Portal'
247-
type = b'3'
248+
type = message_types[b'3'[0]]
248249
__slots__ = ()
249250
CloseCompleteMessage = Message.__new__(CloseComplete)
250251
CloseComplete.SingleInstance = CloseCompleteMessage
251252

252253
class Suspension(EmptyMessage):
253254
'Portal was suspended, more tuples for reading'
254-
type = b's'
255+
type = message_types[b's'[0]]
255256
__slots__ = ()
256257
SuspensionMessage = Message.__new__(Suspension)
257258
Suspension.SingleInstance = SuspensionMessage
258259

259260
class Ready(Message):
260261
'Ready for new query'
261-
type = b'Z'
262+
type = message_types[b'Z'[0]]
262263
__slots__ = ('xact_state',)
263264

264265
def __init__(self, data):
@@ -269,7 +270,7 @@ def serialize(self):
269270

270271
class Notice(Message, dict):
271272
"""Notification message"""
272-
type = b'N'
273+
type = message_types[b'N'[0]]
273274
_dtm = {
274275
b'S' : 'severity',
275276
b'C' : 'code',
@@ -322,12 +323,12 @@ def parse(typ, data):
322323

323324
class Error(Notice):
324325
"""Incoming error"""
325-
type = b'E'
326+
type = message_types[b'E'[0]]
326327
__slots__ = ()
327328

328329
class FunctionResult(Message):
329330
"""Function result value"""
330-
type = b'V'
331+
type = message_types[b'V'[0]]
331332
__slots__ = ('result',)
332333

333334
def __init__(self, datum):
@@ -353,7 +354,7 @@ def parse(typ, data):
353354

354355
class AttributeTypes(TupleMessage):
355356
"""Tuple attribute types"""
356-
type = b't'
357+
type = message_types[b't'[0]]
357358
__slots__ = ()
358359

359360
def serialize(self):
@@ -369,7 +370,7 @@ def parse(typ, data):
369370

370371
class TupleDescriptor(TupleMessage):
371372
"""Tuple description"""
372-
type = b'T'
373+
type = message_types[b'T'[0]]
373374
struct = Struct("!LhLhlh")
374375
__slots__ = ()
375376

@@ -401,7 +402,7 @@ def parse(typ, data):
401402

402403
class Tuple(TupleMessage):
403404
"""Incoming tuple"""
404-
type = b'D'
405+
type = message_types[b'D'[0]]
405406
__slots__ = ()
406407

407408
def serialize(self):
@@ -437,7 +438,7 @@ def parse(typ, data):
437438

438439
class KillInformation(Message):
439440
'Backend cancellation information'
440-
type = b'K'
441+
type = message_types[b'K'[0]]
441442
struct = Struct("!LL")
442443
__slots__ = ('pid', 'key')
443444

@@ -562,7 +563,7 @@ def parse(typ, data):
562563

563564
class Authentication(Message):
564565
"""Authentication(request, salt)"""
565-
type = b'R'
566+
type = message_types[b'R'[0]]
566567
__slots__ = ('request', 'salt')
567568

568569
def __init__(self, request, salt):
@@ -578,38 +579,38 @@ def parse(typ, data):
578579

579580
class Password(StringMessage):
580581
'Password supplement'
581-
type = b'p'
582+
type = message_types[b'p'[0]]
582583
__slots__ = ('data',)
583584

584585
class Disconnect(EmptyMessage):
585586
'Close the connection'
586-
type = b'X'
587+
type = message_types[b'X'[0]]
587588
__slots__ = ()
588589
DisconnectMessage = Message.__new__(Disconnect)
589590
Disconnect.SingleInstance = DisconnectMessage
590591

591592
class Flush(EmptyMessage):
592593
'Flush'
593-
type = b'H'
594+
type = message_types[b'H'[0]]
594595
__slots__ = ()
595596
FlushMessage = Message.__new__(Flush)
596597
Flush.SingleInstance = FlushMessage
597598

598599
class Synchronize(EmptyMessage):
599600
'Synchronize'
600-
type = b'S'
601+
type = message_types[b'S'[0]]
601602
__slots__ = ()
602603
SynchronizeMessage = Message.__new__(Synchronize)
603604
Synchronize.SingleInstance = SynchronizeMessage
604605

605606
class Query(StringMessage):
606607
"""Execute the query with the given arguments"""
607-
type = b'Q'
608+
type = message_types[b'Q'[0]]
608609
__slots__ = ('data',)
609610

610611
class Parse(Message):
611612
"""Parse a query with the specified argument types"""
612-
type = b'P'
613+
type = message_types[b'P'[0]]
613614
__slots__ = ('name', 'statement', 'argtypes')
614615

615616
def __init__(self, name, statement, argtypes):
@@ -645,7 +646,7 @@ class Bind(Message):
645646
rformats, # Result formats; Sequence of BinaryFormat or StringFormat.
646647
)
647648
"""
648-
type = b'B'
649+
type = message_types[b'B'[0]]
649650
__slots__ = ('name', 'statement', 'aformats', 'arguments', 'rformats')
650651

651652
def __init__(self, name, statement, aformats, arguments, rformats):
@@ -704,7 +705,7 @@ def parse(typ, message_data):
704705

705706
class Execute(Message):
706707
"""Fetch results from the specified Portal"""
707-
type = b'E'
708+
type = message_types[b'E'[0]]
708709
__slots__ = ('name', 'max')
709710

710711
def __init__(self, name, max = 0):
@@ -721,7 +722,7 @@ def parse(typ, data):
721722

722723
class Describe(StringMessage):
723724
"""Describe a Portal or Prepared Statement"""
724-
type = b'D'
725+
type = message_types[b'D'[0]]
725726
__slots__ = ('data',)
726727

727728
def serialize(self):
@@ -738,16 +739,16 @@ def parse(typ, data):
738739
return super().parse(data[1:])
739740

740741
class DescribeStatement(Describe):
741-
subtype = b'S'
742+
subtype = message_types[b'S'[0]]
742743
__slots__ = ('data',)
743744

744745
class DescribePortal(Describe):
745-
subtype = b'P'
746+
subtype = message_types[b'P'[0]]
746747
__slots__ = ('data',)
747748

748749
class Close(StringMessage):
749750
"""Generic Close"""
750-
type = b'C'
751+
type = message_types[b'C'[0]]
751752
__slots__ = ()
752753

753754
def serialize(self):
@@ -765,17 +766,17 @@ def parse(typ, data):
765766

766767
class CloseStatement(Close):
767768
"""Close the specified Statement"""
768-
subtype = b'S'
769+
subtype = message_types[b'S'[0]]
769770
__slots__ = ()
770771

771772
class ClosePortal(Close):
772773
"""Close the specified Portal"""
773-
subtype = b'P'
774+
subtype = message_types[b'P'[0]]
774775
__slots__ = ()
775776

776777
class Function(Message):
777778
"""Execute the specified function with the given arguments"""
778-
type = b'F'
779+
type = message_types[b'F'[0]]
779780
__slots__ = ('oid', 'aformats', 'arguments', 'rformat')
780781

781782
def __init__(self, oid, aformats, args, rformat = StringFormat):
@@ -847,16 +848,16 @@ def parse(typ, data):
847848

848849
class CopyToBegin(CopyBegin):
849850
"""Begin copying to"""
850-
type = b'H'
851+
type = message_types[b'H'[0]]
851852
__slots__ = ('format', 'formats')
852853

853854
class CopyFromBegin(CopyBegin):
854855
"""Begin copying from"""
855-
type = b'G'
856+
type = message_types[b'G'[0]]
856857
__slots__ = ('format', 'formats')
857858

858859
class CopyData(Message):
859-
type = b'd'
860+
type = message_types[b'd'[0]]
860861
__slots__ = ('data',)
861862

862863
def __init__(self, data):
@@ -870,11 +871,11 @@ def parse(typ, data):
870871
return typ(data)
871872

872873
class CopyFail(StringMessage):
873-
type = b'f'
874+
type = message_types[b'f'[0]]
874875
__slots__ = ('data',)
875876

876877
class CopyDone(EmptyMessage):
877-
type = b'c'
878+
type = message_types[b'c'[0]]
878879
__slots__ = ('data',)
879880
CopyDoneMessage = Message.__new__(CopyDone)
880881
CopyDone.SingleInstance = CopyDoneMessage

0 commit comments

Comments
 (0)