Skip to content

Commit df018e8

Browse files
ajdavisajdavis
committed
PYTHON-1332 - Gossip $clusterTime
1 parent f0b847a commit df018e8

22 files changed

+459
-160
lines changed

‎pymongo/_cmessagemodule.c‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -388,6 +388,7 @@ static PyObject* _cbson_query_message(PyObject* self, PyObject* args) {
388388
struct module_state *state = GETSTATE(self);
389389

390390
int request_id = rand();
391+
PyObject* cluster_time = NULL;
391392
unsigned int flags;
392393
char* collection_name = NULL;
393394
int collection_name_length;
@@ -429,6 +430,34 @@ static PyObject* _cbson_query_message(PyObject* self, PyObject* args) {
429430
PyErr_NoMemory();
430431
return NULL;
431432
}
433+
434+
/* Pop $clusterTime from dict and write it at the end, avoiding an error
435+
* from the $-prefix and check_keys.
436+
*
437+
* If "dict" is a defaultdict we don't want to call PyMapping_GetItemString
438+
* on it. That would **create** an _id where one didn't previously exist
439+
* (PYTHON-871).
440+
*/
441+
if (PyDict_Check(query)) {
442+
cluster_time = PyDict_GetItemString(query, "$clusterTime");
443+
if (cluster_time) {
444+
/* PyDict_GetItemString returns a borrowed reference. */
445+
Py_INCREF(cluster_time);
446+
if (-1 == PyMapping_DelItemString(query, "$clusterTime")) {
447+
destroy_codec_options(&options);
448+
PyMem_Free(collection_name);
449+
return NULL;
450+
}
451+
}
452+
} else if (PyMapping_HasKeyString(query, "$clusterTime")) {
453+
cluster_time = PyMapping_GetItemString(query, "$clusterTime");
454+
if (!cluster_time
455+
|| -1 == PyMapping_DelItemString(query, "$clusterTime")) {
456+
destroy_codec_options(&options);
457+
PyMem_Free(collection_name);
458+
return NULL;
459+
}
460+
}
432461
if (!buffer_write_int32(buffer, (int32_t)request_id) ||
433462
!buffer_write_bytes(buffer, "\x00\x00\x00\x00\xd4\x07\x00\x00", 8) ||
434463
!buffer_write_int32(buffer, (int32_t)flags) ||
@@ -439,6 +468,7 @@ static PyObject* _cbson_query_message(PyObject* self, PyObject* args) {
439468
destroy_codec_options(&options);
440469
buffer_free(buffer);
441470
PyMem_Free(collection_name);
471+
Py_XDECREF(cluster_time);
442472
return NULL;
443473
}
444474

@@ -447,8 +477,49 @@ static PyObject* _cbson_query_message(PyObject* self, PyObject* args) {
447477
destroy_codec_options(&options);
448478
buffer_free(buffer);
449479
PyMem_Free(collection_name);
480+
Py_XDECREF(cluster_time);
450481
return NULL;
451482
}
483+
484+
/* back up a byte and write $clusterTime */
485+
if (cluster_time) {
486+
int length;
487+
char zero = 0;
488+
489+
buffer_update_position(buffer, buffer_get_position(buffer) - 1);
490+
if (!write_pair(state->_cbson, buffer, "$clusterTime", 12, cluster_time,
491+
0, &options, 1)) {
492+
destroy_codec_options(&options);
493+
buffer_free(buffer);
494+
PyMem_Free(collection_name);
495+
Py_DECREF(cluster_time);
496+
return NULL;
497+
}
498+
499+
if (!buffer_write_bytes(buffer, &zero, 1)) {
500+
destroy_codec_options(&options);
501+
buffer_free(buffer);
502+
PyMem_Free(collection_name);
503+
Py_DECREF(cluster_time);
504+
return NULL;
505+
}
506+
507+
length = buffer_get_position(buffer) - begin;
508+
buffer_write_int32_at_position(buffer, begin, (int32_t)length);
509+
510+
/* undo popping $clusterTime */
511+
if (-1 == PyMapping_SetItemString(
512+
query, "$clusterTime", cluster_time)) {
513+
destroy_codec_options(&options);
514+
buffer_free(buffer);
515+
PyMem_Free(collection_name);
516+
Py_DECREF(cluster_time);
517+
return NULL;
518+
}
519+
520+
Py_DECREF(cluster_time);
521+
}
522+
452523
max_size = buffer_get_position(buffer) - begin;
453524

454525
if (field_selector != Py_None) {

‎pymongo/bulk.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -307,7 +307,8 @@ def execute_command(self, sock_info, generator, write_concern, session):
307307
}
308308
op_id = _randint()
309309
db_name = self.collection.database.name
310-
listeners = self.collection.database.client._event_listeners
310+
client = self.collection.database.client
311+
listeners = client._event_listeners
311312

312313
with self.collection.database.client._tmp_session(session) as s:
313314
# sock_info.command checks auth, but we use sock_info.write_command.
@@ -321,6 +322,7 @@ def execute_command(self, sock_info, generator, write_concern, session):
321322
cmd['bypassDocumentValidation'] = True
322323
if s:
323324
cmd['lsid'] = s._use_lsid()
325+
client._send_cluster_time(cmd)
324326
bwc = _BulkWriteContext(db_name, cmd, sock_info, op_id,
325327
listeners, s)
326328

@@ -329,6 +331,9 @@ def execute_command(self, sock_info, generator, write_concern, session):
329331
run.ops, True, self.collection.codec_options, bwc)
330332

331333
_merge_command(run, full_result, results)
334+
last_result = results[-1][1]
335+
client._receive_cluster_time(last_result)
336+
332337
# We're supposed to continue if errors are
333338
# at the write concern level (e.g. wtimeout)
334339
if self.ordered and full_result['writeErrors']:

‎pymongo/collection.py‎

Lines changed: 16 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
message)
3232
from pymongo.bulk import BulkOperationBuilder, _Bulk
3333
from pymongo.command_cursor import CommandCursor, RawBatchCommandCursor
34+
from pymongo.common import ORDERED_TYPES
3435
from pymongo.collation import validate_collation_or_none
3536
from pymongo.change_stream import ChangeStream
3637
from pymongo.cursor import Cursor, RawBatchCursor
@@ -47,12 +48,6 @@
4748
UpdateResult)
4849
from pymongo.write_concern import WriteConcern
4950

50-
try:
51-
from collections import OrderedDict
52-
_ORDERED_TYPES = (SON, OrderedDict)
53-
except ImportError:
54-
_ORDERED_TYPES = (SON,)
55-
5651
_NO_OBJ_ERROR = "No matching object found"
5752
_UJOIN = u"%s.%s"
5853

@@ -243,7 +238,8 @@ def _command(self, sock_info, command, slave_ok=False,
243238
write_concern=write_concern,
244239
parse_write_concern_error=parse_write_concern_error,
245240
collation=collation,
246-
session=s)
241+
session=s,
242+
client=self.__database.client)
247243

248244
def __create(self, options, collation, session):
249245
"""Sends a create command with the given options.
@@ -573,7 +569,8 @@ def _insert_one(
573569
command,
574570
codec_options=self.__write_response_codec_options,
575571
check_keys=check_keys,
576-
session=s)
572+
session=s,
573+
client=self.__database.client)
577574
_check_write_command_response([(0, result)])
578575
else:
579576
# Legacy OP_INSERT.
@@ -811,7 +808,9 @@ def _update(self, sock_info, criteria, document, upsert=False,
811808
self.__database.name,
812809
command,
813810
codec_options=self.__write_response_codec_options,
814-
session=s).copy()
811+
session=s,
812+
client=self.__database.client).copy()
813+
815814
_check_write_command_response([(0, result)])
816815
# Add the updatedExisting field for compatibility.
817816
if result.get('n') and 'upserted' not in result:
@@ -1089,7 +1088,8 @@ def _delete(
10891088
self.__database.name,
10901089
command,
10911090
codec_options=self.__write_response_codec_options,
1092-
session=s)
1091+
session=s,
1092+
client=self.__database.client)
10931093
_check_write_command_response([(0, result)])
10941094
return result
10951095
else:
@@ -2043,7 +2043,8 @@ def _aggregate(self, pipeline, cursor_class, first_batch_size, session,
20432043
parse_write_concern_error=dollar_out,
20442044
read_concern=read_concern,
20452045
collation=collation,
2046-
session=session)
2046+
session=session,
2047+
client=self.__database.client)
20472048

20482049
if "cursor" in result:
20492050
cursor = result["cursor"]
@@ -2349,8 +2350,9 @@ def rename(self, new_name, session=None, **kwargs):
23492350
if sock_info.max_wire_version >= 5 and self.write_concern:
23502351
cmd['writeConcern'] = self.write_concern.document
23512352
cmd.update(kwargs)
2352-
sock_info.command('admin', cmd, parse_write_concern_error=True,
2353-
session=s)
2353+
return sock_info.command(
2354+
'admin', cmd, parse_write_concern_error=True,
2355+
session=s, client=self.__database.client)
23542356

23552357
def distinct(self, key, filter=None, session=None, **kwargs):
23562358
"""Get a list of distinct values for `key` among all documents
@@ -2974,7 +2976,7 @@ def find_and_modify(self, query={}, update=None,
29742976
kwargs['sort'] = helpers._index_document(sort)
29752977
# Accept OrderedDict, SON, and dict with len == 1 so we
29762978
# don't break existing code already using find_and_modify.
2977-
elif (isinstance(sort, _ORDERED_TYPES) or
2979+
elif (isinstance(sort, ORDERED_TYPES) or
29782980
isinstance(sort, dict) and len(sort) == 1):
29792981
warnings.warn("Passing mapping types for `sort` is deprecated,"
29802982
" use a list of (key, direction) pairs instead",

‎pymongo/command_cursor.py‎

Lines changed: 21 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -129,6 +129,10 @@ def kill():
129129
client = self.__collection.database.client
130130
listeners = client._event_listeners
131131
publish = listeners.enabled_for_commands
132+
start = datetime.datetime.now()
133+
134+
def duration(): return datetime.datetime.now() - start
135+
132136
try:
133137
response = client._send_message_with_response(
134138
operation, address=self.__address)
@@ -140,27 +144,24 @@ def kill():
140144
kill()
141145
raise
142146

143-
cmd_duration = response.duration
144147
rqst_id = response.request_id
145148
from_command = response.from_command
146149
reply = response.data
147150

148-
if publish:
149-
start = datetime.datetime.now()
150151
try:
151152
docs = self._unpack_response(reply,
152153
self.__id,
153154
self.__collection.codec_options)
154155
if from_command:
156+
client._receive_cluster_time(docs[0])
155157
helpers._check_command_response(docs[0])
156158

157159
except OperationFailure as exc:
158160
kill()
159161

160162
if publish:
161-
duration = (datetime.datetime.now() - start) + cmd_duration
162163
listeners.publish_command_failure(
163-
duration, exc.details, "getMore", rqst_id, self.__address)
164+
duration(), exc.details, "getMore", rqst_id, self.__address)
164165

165166
raise
166167
except NotMasterError as exc:
@@ -169,37 +170,38 @@ def kill():
169170
kill()
170171

171172
if publish:
172-
duration = (datetime.datetime.now() - start) + cmd_duration
173173
listeners.publish_command_failure(
174-
duration, exc.details, "getMore", rqst_id, self.__address)
174+
duration(), exc.details, "getMore", rqst_id, self.__address)
175175

176176
client._reset_server_and_request_check(self.address)
177177
raise
178178
except Exception as exc:
179179
if publish:
180-
duration = (datetime.datetime.now() - start) + cmd_duration
181180
listeners.publish_command_failure(
182-
duration, _convert_exception(exc), "getMore", rqst_id,
181+
duration(), _convert_exception(exc), "getMore", rqst_id,
183182
self.__address)
184183
raise
185184

186185
if from_command:
187186
cursor = docs[0]['cursor']
188187
documents = cursor['nextBatch']
189188
self.__id = cursor['id']
189+
if publish:
190+
listeners.publish_command_success(
191+
duration(), docs[0], "getMore", rqst_id,
192+
self.__address)
190193
else:
191194
documents = docs
192195
self.__id = reply.cursor_id
193196

194-
if publish:
195-
duration = (datetime.datetime.now() - start) + cmd_duration
196-
# Must publish in getMore command response format.
197-
res = {"cursor": {"id": self.__id,
198-
"ns": self.__collection.full_name,
199-
"nextBatch": documents},
200-
"ok": 1}
201-
listeners.publish_command_success(
202-
duration, res, "getMore", rqst_id, self.__address)
197+
if publish:
198+
# Must publish in getMore command response format.
199+
res = {"cursor": {"id": self.__id,
200+
"ns": self.__collection.full_name,
201+
"nextBatch": documents},
202+
"ok": 1}
203+
listeners.publish_command_success(
204+
duration(), res, "getMore", rqst_id, self.__address)
203205

204206
if self.__id == 0:
205207
kill()
@@ -227,6 +229,7 @@ def _refresh(self):
227229
self.__id,
228230
self.__collection.codec_options,
229231
self.__session,
232+
self.__collection.database.client,
230233
self.__max_await_time_ms))
231234
else: # Cursor id is zero nothing else to return
232235
self.__killed = True

‎pymongo/common.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
import datetime
2020
import warnings
2121

22+
from bson import SON
2223
from bson.binary import (STANDARD, PYTHON_LEGACY,
2324
JAVA_LEGACY, CSHARP_LEGACY)
2425
from bson.codec_options import CodecOptions
@@ -32,6 +33,13 @@
3233
from pymongo.ssl_support import validate_cert_reqs
3334
from pymongo.write_concern import WriteConcern
3435

36+
try:
37+
from collections import OrderedDict
38+
ORDERED_TYPES = (SON, OrderedDict)
39+
except ImportError:
40+
ORDERED_TYPES = (SON,)
41+
42+
3543
# Defaults until we connect to a server and get updated limits.
3644
MAX_BSON_SIZE = 16 * (1024 ** 2)
3745
MAX_MESSAGE_SIZE = 2 * MAX_BSON_SIZE

0 commit comments

Comments
 (0)