Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def setUpClass(cls):
from pymongo import read_preferences as mongo_rp
cls.db = mongoengine.connect(
db=cls.db_name,
host=cls.db_host,
port=cls.db_port,
# PyMongo>=2.1 requires an explicit read_preference.
read_preference=mongo_rp.ReadPreference.PRIMARY,
# PyMongo>=2.1 has a 20s timeout, use 100ms instead
serverselectiontimeoutms=cls.server_timeout_ms,
)
.. versionadded:: 0.1.2
.. versionchanged:: 0.6 Now chainable
.. versionchanged:: 0.9 Can provide specific fields to reload
"""
max_depth = 1
if fields and isinstance(fields[0], int):
max_depth = fields[0]
fields = fields[1:]
elif 'max_depth' in kwargs:
max_depth = kwargs['max_depth']
if self.pk is None:
raise self.DoesNotExist('Document does not exist')
obj = self._qs.read_preference(ReadPreference.PRIMARY).filter(
**self._object_key).only(*fields).limit(
1).select_related(max_depth=max_depth)
if obj:
obj = obj[0]
else:
raise self.DoesNotExist('Document does not exist')
for field in obj._data:
if not fields or field in fields:
try:
setattr(self, field, self._reload(field, obj[field]))
except (KeyError, AttributeError):
try:
# If field is a special field, e.g. items is stored as _reserved_items,
# an KeyError is thrown. So try to retrieve the field from _data
def _slaveok_for_server(self, read_preference, server, session,
exhaust=False):
assert read_preference is not None, "read_preference must not be None"
# Get a socket for a server matching the read preference, and yield
# sock_info, slave_ok. Server Selection Spec: "slaveOK must be sent to
# mongods with topology type Single. If the server type is Mongos,
# follow the rules for passing read preference to mongos, even for
# topology type Single."
# Thread safe: if the type is single it cannot change.
topology = self._get_topology()
single = topology.description.topology_type == TOPOLOGY_TYPE.Single
with self._get_socket(server, session, exhaust=exhaust) as sock_info:
slave_ok = (single and not sock_info.is_mongos) or (
read_preference != ReadPreference.PRIMARY)
yield sock_info, slave_ok
def _command(self, sock_info, command, slave_ok=False, value=1, check=True,
allowable_errors=None, read_preference=ReadPreference.PRIMARY,
codec_options=DEFAULT_CODEC_OPTIONS,
write_concern=None,
parse_write_concern_error=False, **kwargs):
"""Internal command helper."""
if isinstance(command, string_type):
command = SON([(command, value)])
if sock_info.max_wire_version >= 5 and write_concern:
command['writeConcern'] = write_concern.document
command.update(kwargs)
return sock_info.command(
self.__name,
command,
slave_ok,
# Configure the connection
self.configure()
if self.config['replica_set'] is not None:
connection = MongoReplicaSetClient(
self.config['uri'],
replicaSet=self.config['replica_set'],
w=self.config['write_concern'],
fsync=self.config['fsync'],
read_preference=ReadPreference.PRIMARY_PREFERRED)
else:
# Connecting to a stand alone MongoDB
connection = MongoClient(
self.config['uri'],
fsync=self.config['fsync'],
read_preference=ReadPreference.PRIMARY)
# Set up the database
self.database = connection[self.config['database']]
self.collections = {'default': self.database[self.config['collection']]}
self.logger.info(u'Connected to MongoDB {0}, using "{1}"'.format(
self.config['uri'],
self.config['database']))
# Get the duplicate on key option
if self.config['stop_on_duplicate']:
tmpValue = self.config['stop_on_duplicate']
if tmpValue < 0:
msg = (
u'Negative values are not allowed for'
if not isinstance(indexes, list):
raise TypeError("indexes must be a list")
names = []
def gen_indexes():
for index in indexes:
if not isinstance(index, IndexModel):
raise TypeError("%r is not an instance of "
"pymongo.operations.IndexModel" % (index,))
document = index.document
names.append(document["name"])
yield document
cmd = SON([('createIndexes', self.name),
('indexes', list(gen_indexes()))])
with self._socket_for_writes() as sock_info:
self._command(
sock_info, cmd, read_preference=ReadPreference.PRIMARY,
codec_options=_UNICODE_REPLACE_CODEC_OPTIONS,
write_concern=self.write_concern,
parse_write_concern_error=True)
return names
def mongo_connect(self):
'''Connects to mongo'''
try:
tools.config['mongodb_replicaset'] = tools.config.get(
'mongodb_replicaset', False
)
mongo_client = MongoClient
kwargs = {}
if tools.config['mongodb_replicaset']:
kwargs.update({'replicaSet': tools.config['mongodb_replicaset'],
'read_preference': ReadPreference.PRIMARY_PREFERRED})
mongo_client = MongoReplicaSetClient
connection = mongo_client(self.uri, **kwargs)
except Exception, e:
raise except_orm('MongoDB connection error', e)
return connection
def _end_sessions(self, session_ids):
"""Send endSessions command(s) with the given session ids."""
try:
# Use SocketInfo.command directly to avoid implicitly creating
# another session.
with self._socket_for_reads(
ReadPreference.PRIMARY_PREFERRED,
None) as (sock_info, slave_ok):
if not sock_info.supports_sessions:
return
for i in range(0, len(session_ids), common._MAX_END_SESSIONS):
spec = SON([('endSessions',
session_ids[i:i + common._MAX_END_SESSIONS])])
sock_info.command(
'admin', spec, slave_ok=slave_ok, client=self)
except PyMongoError:
# Drivers MUST ignore any errors returned by the endSessions
# command.
pass