def test_mongo_client_with_ssl_after_2_1():
if _version._lt('2.1'):
raise SkipTest("This test requires version 2.1 or later.")
class SSLMongo(Mongo):
config_host = 'localhost'
config_port = 27017
config_ssl = True
config_mongo_client = ({'serverSelectionTimeoutMS': 300} if
_version._gte('3.0') else {})
class SomeDoc(Document):
config_database = database_name()
config_collection = 'ssl_collection'
name = 'n'
try:
SomeDoc.insert
except:
raise
try:
import socket
socket.setdefaulttimeout(3)
with SSLMongo:
SomeDoc.insert({SomeDoc.name:'foobar'})
ok_(SomeDoc.find({SomeDoc.name:'foobar'}))
except ConnectionFailure as err:
raise SkipTest("SSL may not be enabled on mongodb server: %r" % err)
def _new_connection(cls):
""" Return a new connection to this class' database. """
kwargs = cls._connection_info()
kwargs.update({
'max_pool_size': cls.config_max_pool_size,
'auto_start_request': cls.config_auto_start_request,
'use_greenlets': cls.config_use_greenlets,
'tz_aware': cls.config_tz_aware,
'w': cls.config_write_concern,
'ssl': cls.config_ssl,
})
if _version._gte('2.1.0') and _version._lt('2.2.0'):
# This causes an error for the 2.1.x versions of Pymongo, so we
# remove it
kwargs.pop('auto_start_request')
kwargs.pop('use_greenlets')
elif _version._gte('3.0.0'):
kwargs.pop('auto_start_request')
kwargs.pop('max_pool_size')
kwargs.pop('use_greenlets')
if cls.config_replica:
kwargs['replicaSet'] = cls.config_replica
logging.getLogger(__name__).info("Creating new MongoDB connection "
"to '{}:{}' replica: {}".format(cls.config_host,
cls.config_port, cls.config_replica))
else:
logging.getLogger(__name__).info("Creating new MongoDB connection "
"to '{}:{}'".format(cls.config_host, cls.config_port))
return cls.config_connection_cls(**kwargs)
class Page(Document):
""" Document class used by :class:`Array`. """
size = 's' # Number of entries in this page
""" Number of entries currently in this page. """
entries = 'e' # Array of entries
""" Array of entries. """
_opts = {'safe': True} if _version._lt('3.0.0') else {}
def reconnect(cls):
""" Replace the current connection with a new connection. """
logging.getLogger(__name__).info("Reloading '{}'"
.format(cls.__name__))
if (cls._connection and _version._lt('3.0.0')):
cls._connection.disconnect()
cls._connection = cls._new_connection()
def next(self):
if six.PY3 and _version._gte('3'):
doc = super().next()
elif six.PY3 and _version._lt('3'):
doc = super().__next__()
else:
doc = super(Cursor, self).next()
doc = self._doc_cls(doc)
return doc
def start(cls):
""" Public function for manually starting a session/context. Use
carefully!
"""
if cls in Mongo.contexts:
raise NestedConnection("Do not nest a connection within itself, it "
"may cause undefined behavior.")
if (pyconfig.get('humbledb.allow_explicit_request', True)
and _version._lt('3.0.0')):
cls.connection.start_request()
Mongo.contexts.append(cls)
def __init__(self, index, cache_for=(60 * 60 * 24), background=True,
**kwargs):
self.index = index
# Merge kwargs
kwargs['cache_for'] = cache_for
kwargs['background'] = background
if _version._lt('2.3') and 'cache_for' in kwargs:
kwargs['ttl'] = kwargs.pop('cache_for')
self.kwargs = kwargs
def end(cls):
""" Public function for manually closing a session/context. Should be
idempotent. This must always be called after :meth:`Mongo.start`
to ensure the socket is returned to the connection pool.
"""
if (pyconfig.get('humbledb.allow_explicit_request', True)
and _version._lt('3.0.0')):
cls.connection.end_request()
try:
Mongo.contexts.pop()
except (IndexError, AttributeError):
pass
def length(self):
""" Return the total number of items in this array. """
# This is implemented rather than __len__ because it incurs a query,
# and we don't want to query transparently
Page = self._page
if _version._lt('3.0.0'):
cursor = Page.find({'_id': self._id_regex}, fields={Page.size:
1, '_id': 0})
else:
cursor = Page.find({'_id': self._id_regex}, {Page.size: 1, '_id':
0})
return sum(p.size for p in cursor)
def test_replica_works_for_versions_after_2_4():
if _version._lt('2.4'):
raise SkipTest
with mock.patch('pymongo.MongoReplicaSetClient') as replica:
class Replica(Mongo):
config_host = 'localhost'
config_port = 27017
config_replica = 'test'
with Replica:
pass
replica.assert_called_once()
def test_replica_works_for_versions_between_2_1_and_2_4():
if _version._lt('2.1') or _version._gte('2.4'):
raise SkipTest
with mock.patch('pymongo.ReplicaSetConnection') as replica:
class Replica(Mongo):
config_host = 'localhost'
config_port = 27017
config_replica = 'test'
with Replica:
pass
replica.assert_called_once()
def test_mongo_uri_with_database():
if _version._lt('2.6.0'):
raise SkipTest("Needs version 2.6.0 or later")
host = pyconfig.get('humbledb.test.db.host', 'localhost')
port = pyconfig.get('humbledb.test.db.port', 27017)
uri = 'mongodb://{}:{}/{}'.format(host, port, database_name())
class DBuri(Mongo):
config_uri = uri
with DBuri:
eq_(DBuri.database.name, database_name())
eq_(Mongo.context.database.name, database_name())
def database(cls):
"""
Return the default database for this connection.
.. versionadded:: 5.3.0
.. note:: This requires ``pymongo >= 2.6.0``.
"""
if _version._lt('2.6.0'):
return None
try:
return cls.connection.get_default_database()
except pymongo.errors.ConfigurationError:
return None
def test_replica_works_for_versions_after_2_4():
if _version._lt('2.4'):
raise SkipTest
with mock.patch('pymongo.MongoReplicaSetClient') as replica:
class Replica(Mongo):
config_host = 'localhost'
config_port = 27017
config_replica = 'test'
with Replica:
pass
replica.assert_called_once()
def database(cls):
"""
Return the default database for this connection.
.. versionadded:: 5.3.0
.. note:: This requires ``pymongo >= 2.6.0``.
"""
if _version._lt('2.6.0'):
return None
try:
return cls.connection.get_default_database()
except pymongo.errors.ConfigurationError:
return None
def _new_connection(cls):
""" Return a new connection to this class' database. """
kwargs = cls._connection_info()
kwargs.update({
'max_pool_size': cls.config_max_pool_size,
'auto_start_request': cls.config_auto_start_request,
'use_greenlets': cls.config_use_greenlets,
'tz_aware': cls.config_tz_aware,
'w': cls.config_write_concern,
'ssl': cls.config_ssl,
})
kwargs.update(cls.config_mongo_client)
if _version._gte('2.1.0') and _version._lt('2.2.0'):
# This causes an error for the 2.1.x versions of Pymongo, so we
# remove it
kwargs.pop('auto_start_request')
kwargs.pop('use_greenlets')
if _version._gte('3.0.0'):
# Handle removed keywords
kwargs.pop('use_greenlets')
kwargs.pop('auto_start_request')
# Handle changed keywords
kwargs['maxPoolSize'] = kwargs.pop('max_pool_size')
#extra timeout stuff because this doesnt allow any other parameters
kwargs['socketTimeoutMS'] = 10000
kwargs['connectTimeoutMS'] = 200
kwargs['serverSelectionTimeoutMS'] = 2000
kwargs['waitQueueTimeoutMS'] = 2000
# Handle other 3.0 stuff
if kwargs.get('ssl') and ssl:
kwargs.setdefault('ssl_cert_reqs', ssl.CERT_NONE)
if cls.config_replica:
kwargs['replicaSet'] = cls.config_replica
logging.getLogger(__name__).info("Creating new MongoDB connection "
"to '{}:{}' replica: {}".format(cls.config_host,
cls.config_port, cls.config_replica))
else:
logging.getLogger(__name__).info("Creating new MongoDB connection "
"to '{}:{}'".format(cls.config_host, cls.config_port))
return cls.config_connection_cls(**kwargs)
def _preallocate(cls, event, stamp):
"""
Preallocate a new document for `event` during the period containing
`stamp`.
:param event: Event identifier string
:param stamp: A UTC datetime indicating the document period
:type event: str
:type stamp: datetime.datetime
"""
# Get the time period for this report
period = cls._period(stamp)
# If we already have preallocated for this time period, get out of here
if event in cls._preallocated[period]:
return
# Do a fast check if the document exists
if cls.find({cls._id: cls.record_id(event, stamp)}).limit(1).count():
return
# Get our query and update clauses
query, update = cls._preallocate_query(event, stamp)
try:
_opts = {}
if _version._lt('3.0.0'):
_opts['safe'] = True
# We always want preallocation to be "safe" in order to avoid race
# conditions with the subsequent update
cls.update(query, update, upsert=True, **_opts)
except humbledb.errors.DuplicateKeyError: # pragma: no cover
# Get out of here, we're done
return
# XXX: This will not scale well! However it should be good up to 100k
# or so event identifiers. Maybe a bloom filter would work better?
# Add the event identifier to the list of already preallocated
# documents for this period
cls._preallocated[period].add(event)
# Clean up the old preallocation records
if len(cls._preallocated) > 2:
previous = _relative_period(cls.config_period, period, -2)
cls._preallocated.pop(previous, None)
def _preallocate(cls, event, stamp):
"""
Preallocate a new document for `event` during the period containing
`stamp`.
:param event: Event identifier string
:param stamp: A UTC datetime indicating the document period
:type event: str
:type stamp: datetime.datetime
"""
# Get the time period for this report
period = cls._period(stamp)
# If we already have preallocated for this time period, get out of here
if event in cls._preallocated[period]:
return
# Do a fast check if the document exists
if cls.find({cls._id: cls.record_id(event, stamp)}).limit(1).count():
return
# Get our query and update clauses
query, update = cls._preallocate_query(event, stamp)
try:
_opts = {}
if _version._lt('3.0.0'):
_opts['safe'] = True
# We always want preallocation to be "safe" in order to avoid race
# conditions with the subsequent update
cls.update(query, update, upsert=True, **_opts)
except humbledb.errors.DuplicateKeyError: # pragma: no cover
# Get out of here, we're done
return
# XXX: This will not scale well! However it should be good up to 100k
# or so event identifiers. Maybe a bloom filter would work better?
# Add the event identifier to the list of already preallocated
# documents for this period
cls._preallocated[period].add(event)
# Clean up the old preallocation records
if len(cls._preallocated) > 2:
previous = _relative_period(cls.config_period, period, -2)
cls._preallocated.pop(previous, None)
def test_mongo_uri_database_with_conflict_raises_error():
if _version._lt('2.6.0'):
raise SkipTest("Needs version 2.6.0 or later")
host = pyconfig.get('humbledb.test.db.host', 'localhost')
port = pyconfig.get('humbledb.test.db.port', 27017)
uri = 'mongodb://{}:{}/{}'.format(host, port, database_name())
class DBuri(Mongo):
config_uri = uri
from humbledb import Document
class TestDoc(Document):
config_database = database_name() + '_is_different'
config_collection = 'test'
with DBuri:
TestDoc.find()
def test_auto_increment_errors_with_wrong_db():
if _version._lt('2.6.0'):
raise SkipTest
host = pyconfig.get('humbledb.test.db.host', 'localhost')
port = pyconfig.get('humbledb.test.db.port', 27017)
uri = 'mongodb://{}:{}/{}'.format(host, port, database_name())
class DBuri(Mongo):
config_uri = uri
class MyDoc2(Document):
config_database = database_name()
config_collection = 'test'
auto = 'a', auto_increment(database_name() + '_is_different', SIDECAR,
'MyDoc2')
doc = MyDoc2()
with DBuri:
doc.auto
def record(cls, event, stamp=None, safe=False, count=1):
"""
Record an instance of `event` that happened at `stamp`.
If `safe` is ``True``, then this method will wait for write
acknowledgement from the server. The `safe` keyword has no effect for
`pymongo >= 3.0.0`.
:param event: Event identifier string
:param stamp: Datetime stamp for this event (default: now)
:param safe: Safe write option passed to pymongo
:param count: Number to increment
:type event: str
:type stamp: datetime.datetime
:type safe: bool
:type count: int
"""
if not isinstance(count, six.integer_types):
raise ValueError("'count' must be int or long, got %r instead" %
type(count))
if stamp and not isinstance(stamp, (datetime.datetime, datetime.date)):
raise ValueError(
"'stamp' must be datetime or date, got %r instead" %
type(stamp))
# Get our stamp as UTC time or use the current time
stamp = pytool.time.as_utc(stamp) if stamp else pytool.time.utcnow()
# Do preallocation
cls._attempt_preallocation(event, stamp)
# Get the update query
update = cls._update_query(stamp, count)
# Get our query doc
doc = {'_id': cls.record_id(event, stamp)}
_opts = {}
if _version._lt('3.0.0'):
_opts['safe'] = safe
# Update/upsert the document, hooray
cls.update(doc, update, upsert=True, **_opts)
def test_auto_increment_errors_with_wrong_db():
if _version._lt('2.6.0'):
raise SkipTest
host = pyconfig.get('humbledb.test.db.host', 'localhost')
port = pyconfig.get('humbledb.test.db.port', 27017)
uri = 'mongodb://{}:{}/{}'.format(host, port, database_name())
class DBuri(Mongo):
config_uri = uri
class MyDoc2(Document):
config_database = database_name()
config_collection = 'test'
auto = 'a', auto_increment(database_name() + '_is_different', SIDECAR,
'MyDoc2')
doc = MyDoc2()
with DBuri:
doc.auto
def record(cls, event, stamp=None, safe=False, count=1):
"""
Record an instance of `event` that happened at `stamp`.
If `safe` is ``True``, then this method will wait for write
acknowledgement from the server. The `safe` keyword has no effect for
`pymongo >= 3.0.0`.
:param event: Event identifier string
:param stamp: Datetime stamp for this event (default: now)
:param safe: Safe write option passed to pymongo
:param count: Number to increment
:type event: str
:type stamp: datetime.datetime
:type safe: bool
:type count: int
"""
if not isinstance(count, six.integer_types):
raise ValueError("'count' must be int or long, got %r instead" %
type(count))
if stamp and not isinstance(stamp, (datetime.datetime, datetime.date)):
raise ValueError("'stamp' must be datetime or date, got %r instead"
% type(stamp))
# Get our stamp as UTC time or use the current time
stamp = pytool.time.as_utc(stamp) if stamp else pytool.time.utcnow()
# Do preallocation
cls._attempt_preallocation(event, stamp)
# Get the update query
update = cls._update_query(stamp, count)
# Get our query doc
doc = {'_id': cls.record_id(event, stamp)}
_opts = {}
if _version._lt('3.0.0'):
_opts['safe'] = safe
# Update/upsert the document, hooray
cls.update(doc, update, upsert=True, **_opts)
def test_mongo_client_with_ssl_after_2_1():
if _version._lt('2.1'):
raise SkipTest("This test requires version 2.1 or later.")
class SSLMongo(Mongo):
config_host = 'localhost'
config_port = 27017
config_ssl = True
from humbledb import Document
from humbledb.errors import ConnectionFailure
class SomeDoc(Document):
config_database = database_name()
config_collection = 'ssl_collection'
name = 'n'
try:
with SSLMongo:
SomeDoc.insert({SomeDoc.name:'foobar'})
ok_(SomeDoc.find({SomeDoc.name:'foobar'}))
except ConnectionFailure:
raise SkipTest("SSL may not be enabled on mongodb server.")
def test_mongo_client_with_ssl_after_2_1():
if _version._lt('2.1'):
raise SkipTest("This test requires version 2.1 or later.")
class SSLMongo(Mongo):
config_host = 'localhost'
config_port = 27017
config_ssl = True
from humbledb import Document
from humbledb.errors import ConnectionFailure
class SomeDoc(Document):
config_database = database_name()
config_collection = 'ssl_collection'
name = 'n'
try:
with SSLMongo:
SomeDoc.insert({SomeDoc.name: 'foobar'})
ok_(SomeDoc.find({SomeDoc.name: 'foobar'}))
except ConnectionFailure:
raise SkipTest("SSL may not be enabled on mongodb server.")
def cache_for(val):
# This is a work around for the version changing the cache argument
if _version._lt('2.3'):
return {'ttl': val}
return {'cache_for': val}
def cache_for(val):
# This is a work around for the version changing the cache argument
if _version._lt('2.3'):
return {'ttl': val}
return {'cache_for': val}
def __new__(mcs, name, bases, cls_dict):
""" Return the Mongo class. """
# This ensures that a late-declared class does not inherit an existing
# connection object.
cls_dict['_connection'] = None
# Choose the correct connection class
if cls_dict.get('config_connection_cls', UNSET) is UNSET:
# Are we using a replica?
# XXX: Getting the connection type at class creation time rather
# than connection instantiation time means that disabling
# config_replica (setting to None) at runtime has no effect. I
# doubt anyone would ever do this, but you never know.
_replica = cls_dict.get('config_replica', UNSET)
# Handle attribute descriptors responsibly
if _replica and hasattr(_replica, '__get__'):
try:
_replica = _replica.__get__(None, None)
except:
raise TypeError("'%s.config_replica' appears to be a "
"descriptor and its value could not be "
"retrieved reliably." % name)
# Handle replica set connections
if _replica:
if _version._lt('2.1'):
raise TypeError("Need pymongo.version >= 2.1 for replica "
"sets.")
elif _version._gte('2.4'):
conn = pymongo.MongoReplicaSetClient
else:
conn = pymongo.ReplicaSetConnection
else:
# Get the correct regular connection
if _version._gte('2.4'):
conn = pymongo.MongoClient
else:
conn = pymongo.Connection
# Set our connection type
cls_dict['config_connection_cls'] = conn
# Specially handle base class
if name == 'Mongo' and bases == (object,):
# Create thread local self
cls_dict['_self'] = threading.local()
return type.__new__(mcs, name, bases, cls_dict)
if cls_dict.get('config_uri', UNSET) is UNSET:
# Ensure we have minimum configuration params
if cls_dict.get('config_host', UNSET) is UNSET:
raise TypeError("missing required 'config_host'")
if cls_dict.get('config_port', UNSET) is UNSET:
raise TypeError("missing required 'config_port'")
# Validate if pymongo version supports SSL.
if (cls_dict.get('config_ssl', False) is True and
_version._lt('2.1')):
raise TypeError("Need pymongo.version >= 2.1 to use "
"SSL.")
# Create new class
cls = type.__new__(mcs, name, bases, cls_dict)
# This reload hook uses a closure to access the class
@pyconfig.reload_hook
def _reload():
""" A hook for reloading the connection settings with pyconfig. """
cls.reconnect()
return cls
Tests for Helpers
=================
"""
import pyconfig
from humbledb import _version
from humbledb import Document, Mongo
from humbledb.helpers import auto_increment
from humbledb.errors import DatabaseMismatch, NoConnection
from ..util import DBTest, database_name, eq_, ok_, raises, SkipTest
SIDECAR = 'sidecars'
# The safe= keyword doesn't exist in 3.0
if _version._lt('3.0.0'):
_safe = {'safe': True}
else:
_safe = {}
class MyDoc(Document):
config_database = database_name()
config_collection = 'test'
auto = 'a', auto_increment(database_name(), SIDECAR, 'MyDoc')
class MyFloatCounterDoc(Document):
config_database = database_name()
config_collection = 'float_test'
import mock
import pytool
import pyconfig
import humbledb
from humbledb import Mongo, Document, Embed, _version
from ..util import eq_, ok_, raises, DBTest, database_name
# The safe= keyword doesn't exist in 3.0
if _version._lt('3.0.0'):
_safe = {'safe': True}
else:
_safe = {}
def teardown():
DBTest.connection.drop_database(database_name())
def cache_for(val):
# This is a work around for the version changing the cache argument
if _version._lt('2.3'):
return {'ttl': val}
return {'cache_for': val}
class DocTest(Document):
config_database = database_name()
config_collection = 'test'
def __new__(mcs, name, bases, cls_dict):
""" Return the Mongo class. """
# This ensures that a late-declared class does not inherit an existing
# connection object.
cls_dict['_connection'] = None
# Choose the correct connection class
if cls_dict.get('config_connection_cls', UNSET) is UNSET:
# Are we using a replica?
# XXX: Getting the connection type at class creation time rather
# than connection instantiation time means that disabling
# config_replica (setting to None) at runtime has no effect. I
# doubt anyone would ever do this, but you never know.
_replica = cls_dict.get('config_replica', UNSET)
# Handle attribute descriptors responsibly
if _replica and hasattr(_replica, '__get__'):
try:
_replica = _replica.__get__(None, None)
except:
raise TypeError("'%s.config_replica' appears to be a "
"descriptor and its value could not be "
"retrieved reliably." % name)
# Handle replica set connections
if _replica:
if _version._lt('2.1'):
raise TypeError("Need pymongo.version >= 2.1 for replica "
"sets.")
elif _version._gte('3.0.0'):
conn = pymongo.MongoClient
elif _version._gte('2.4'):
conn = pymongo.MongoReplicaSetClient
else:
conn = pymongo.ReplicaSetConnection
else:
# Get the correct regular connection
if _version._gte('2.4'):
conn = pymongo.MongoClient
else:
conn = pymongo.Connection
# Set our connection type
cls_dict['config_connection_cls'] = conn
# Specially handle base class
if name == 'Mongo' and bases == (object,):
# Create thread local self
cls_dict['_self'] = threading.local()
return type.__new__(mcs, name, bases, cls_dict)
if cls_dict.get('config_uri', UNSET) is UNSET:
# Ensure we have minimum configuration params
if cls_dict.get('config_host', UNSET) is UNSET:
raise TypeError("missing required 'config_host'")
if cls_dict.get('config_port', UNSET) is UNSET:
raise TypeError("missing required 'config_port'")
# Validate if pymongo version supports SSL.
if (cls_dict.get('config_ssl', False) is True and
_version._lt('2.1')):
raise TypeError("Need pymongo.version >= 2.1 to use "
"SSL.")
# Create new class
cls = type.__new__(mcs, name, bases, cls_dict)
# This reload hook uses a closure to access the class
@pyconfig.reload_hook
def _reload():
""" A hook for reloading the connection settings with pyconfig. """
cls.reconnect()
return cls