import logging
import os
import re
import threading
import pymongo
from pymongo.errors import OperationFailure, AutoReconnect
from six import string_types
from ._cache import Cache
from ._config import ENABLE_CACHE
from ._util import indent
from .auth import authenticate, get_auth
from .chunkstore import chunkstore
from .decorators import mongo_retry
from .exceptions import LibraryNotFoundException, ArcticException, QuotaExceededException
from .hooks import get_mongodb_uri
from .store import version_store, bson_store, metadata_store
from .tickstore import tickstore, toplevel
__all__ = ['Arctic', 'VERSION_STORE', 'METADATA_STORE', 'TICK_STORE', 'CHUNK_STORE', 'register_library_type']
# Set default logging handler to avoid "No handler found" warnings.
logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())
# Default Arctic application name: 'arctic'
APPLICATION_NAME = 'arctic'
VERSION_STORE = version_store.VERSION_STORE_TYPE
METADATA_STORE = metadata_store.METADATA_STORE_TYPE
TICK_STORE = tickstore.TICK_STORE_TYPE
CHUNK_STORE = chunkstore.CHUNK_STORE_TYPE
LIBRARY_TYPES = {version_store.VERSION_STORE_TYPE: version_store.VersionStore,
tickstore.TICK_STORE_TYPE: tickstore.TickStore,
toplevel.TICK_STORE_TYPE: toplevel.TopLevelTickStore,
chunkstore.CHUNK_STORE_TYPE: chunkstore.ChunkStore,
bson_store.BSON_STORE_TYPE: bson_store.BSONStore,
metadata_store.METADATA_STORE_TYPE: metadata_store.MetadataStore
}
def register_library_type(name, type_):
"""
Register a Arctic Library Type handler
"""
if name in LIBRARY_TYPES:
raise ArcticException("Library %s already registered as %s" % (name, LIBRARY_TYPES[name]))
LIBRARY_TYPES[name] = type_
class Arctic(object):
"""
The Arctic class is a top-level God object, owner of all arctic_<user> databases
accessible in Mongo.
Each database contains one or more ArcticLibrarys which may have implementation
specific functionality.
Current Mongo Library types:
- arctic.VERSION_STORE - Versioned store for chunked Pandas and numpy objects
(other Python types are pickled)
- arctic.TICK_STORE - Tick specific library. Supports 'snapshots', efficiently
stores updates, not versioned.
- arctic.METADATA_STORE - Stores metadata with timestamps
Arctic and ArcticLibrary are responsible for Connection setup, authentication,
dispatch to the appropriate library implementation, and quotas.
"""
DB_PREFIX = 'arctic'
METADATA_COLL = "ARCTIC"
METADATA_DOC_ID = "ARCTIC_META"
_MAX_CONNS = 4
__conn = None
def __init__(self, mongo_host, app_name=APPLICATION_NAME, allow_secondary=False,
socketTimeoutMS=10 * 60 * 1000, connectTimeoutMS=2 * 1000,
serverSelectionTimeoutMS=30 * 1000, **kwargs):
"""
Constructs a Arctic Datastore.
Note: If mongo_host is a pymongo connection and the process is later forked, the
new pymongo connection may have different parameters.
Parameters:
-----------
mongo_host: A MongoDB hostname, alias or Mongo Connection
app_name: `str` is the name of application used for resolving credentials when
authenticating against the mongo_host.
We will fetch credentials using the authentication hook.
Teams should override this such that different applications don't accidentally
run with privileges to other applications' databases
allow_secondary: `bool` indicates if we allow reads against
secondary members in the cluster. These reads may be
a few seconds behind (but are usually split-second up-to-date).
serverSelectionTimeoutMS: `int` the main tunable used for configuring how long
the pymongo driver will spend on MongoDB cluster discovery. This parameter
takes precedence over connectTimeoutMS: https://jira.mongodb.org/browse/DRIVERS-222
kwargs: 'dict' extra keyword arguments to pass when calling pymongo.MongoClient,
for example ssl parameters.
"""
self._application_name = app_name
self._library_cache = {}
self._allow_secondary = allow_secondary
self._socket_timeout = socketTimeoutMS
self._connect_timeout = connectTimeoutMS
self._server_selection_timeout = serverSelectionTimeoutMS
self._lock = threading.RLock()
self._pid = os.getpid()
self._pymongo_kwargs = kwargs
self._cache = None
if isinstance(mongo_host, string_types):
self._given_instance = False
self.mongo_host = mongo_host
else:
self._given_instance = True
self.__conn = mongo_host
# Workaround for: https://jira.mongodb.org/browse/PYTHON-927
mongo_host.server_info()
self.mongo_host = ",".join(["{}:{}".format(x[0], x[1]) for x in mongo_host.nodes])
self._adminDB = self._conn.admin
self._cache = Cache(self._conn)
@property
@mongo_retry
def _conn(self):
with self._lock:
# We must make sure that no MongoClient instances are used from parent after fork:
# http://api.mongodb.com/python/current/faq.html#using-pymongo-with-multiprocessing
curr_pid = os.getpid()
if curr_pid != self._pid:
if self._given_instance:
logger.warn("Forking process. Arctic was passed a pymongo connection during init, "
"the new pymongo connection may have different parameters.")
self._pid = curr_pid # this line has to precede reset() otherwise we get to eternal recursion
self.reset() # also triggers re-auth
if self.__conn is None:
host = get_mongodb_uri(self.mongo_host)
logger.info("Connecting to mongo: {0} ({1})".format(self.mongo_host, host))
self.__conn = pymongo.MongoClient(host=host,
maxPoolSize=self._MAX_CONNS,
socketTimeoutMS=self._socket_timeout,
connectTimeoutMS=self._connect_timeout,
serverSelectionTimeoutMS=self._server_selection_timeout,
**self._pymongo_kwargs)
self._adminDB = self.__conn.admin
self._cache = Cache(self.__conn)
# Authenticate against admin for the user
auth = get_auth(self.mongo_host, self._application_name, 'admin')
if auth:
authenticate(self._adminDB, auth.user, auth.password)
# Accessing _conn is synchronous. The new PyMongo driver may be lazier than the previous.
# Force a connection.
self.__conn.server_info()
return self.__conn
def reset(self):
logger.debug("Arctic.reset()")
with self._lock:
if self.__conn is not None:
self.__conn.close()
self.__conn = None
for _, l in self._library_cache.items():
if hasattr(l, '_reset') and callable(l._reset):
logger.debug("Library reset() %s" % l)
l._reset() # the existence of _reset() is not guaranteed/enforced, it also triggers re-auth
def __str__(self):
return "<Arctic at %s, connected to %s>" % (hex(id(self)), str(self._conn))
def __repr__(self):
return str(self)
def __getstate__(self):
return {'mongo_host': self.mongo_host,
'app_name': self._application_name,
'allow_secondary': self._allow_secondary,
'socketTimeoutMS': self._socket_timeout,
'connectTimeoutMS': self._connect_timeout,
'serverSelectionTimeoutMS': self._server_selection_timeout}
def __setstate__(self, state):
return Arctic.__init__(self, **state)
def is_caching_enabled(self):
"""
Allows people to enable or disable caching for list_libraries globally.
"""
_ = self._conn # Ensures the connection exists and cache is initialized with it.
return self._cache.is_caching_enabled(ENABLE_CACHE)
def list_libraries(self, newer_than_secs=None):
"""
Returns
-------
list of Arctic library names
"""
return self._list_libraries_cached(newer_than_secs) if self.is_caching_enabled() else self._list_libraries()
@mongo_retry
def _list_libraries(self):
libs = []
for db in self._conn.list_database_names():
if db.startswith(self.DB_PREFIX + '_'):
for coll in self._conn[db].list_collection_names():
if coll.endswith(self.METADATA_COLL):
libs.append(db[len(self.DB_PREFIX) + 1:] + "." + coll[:-1 * len(self.METADATA_COLL) - 1])
elif db == self.DB_PREFIX:
for coll in self._conn[db].list_collection_names():
if coll.endswith(self.METADATA_COLL):
libs.append(coll[:-1 * len(self.METADATA_COLL) - 1])
return libs
# Better to be pessimistic here and not retry.
def _list_libraries_cached(self, newer_than_secs=None):
"""
Returns
-------
List of Arctic library names from a cached collection (global per mongo cluster) in mongo.
Long term list_libraries should have a use_cached argument.
"""
_ = self._conn # Ensures the connection exists and cache is initialized with it.
cache_data = self._cache.get('list_libraries', newer_than_secs)
if not cache_data:
# Try to refresh the cache.
logging.debug("Cache has expired data, fetching from slow path and reloading cache.")
libs = self._list_libraries()
self._cache.set('list_libraries', libs)
return libs
return cache_data
def reload_cache(self):
_ = self._conn # Ensures the connection exists and cache is initialized with it.
self._cache.set('list_libraries', self._list_libraries())
def library_exists(self, library):
"""
Check whether a given library exists.
Parameters
----------
library : `str`
The name of the library. e.g. 'library' or 'user.library'
Returns
-------
`bool`
True if the library with the given name already exists, False otherwise
"""
exists = False
try:
# This forces auth errors, and to fall back to the slower "list_collections"
ArcticLibraryBinding(self, library).get_library_type()
# This will obtain the library, if no exception thrown we have verified its existence
self.get_library(library)
exists = True
except OperationFailure:
exists = library in self.list_libraries()
except LibraryNotFoundException:
pass
return exists
def _sanitize_lib_name(self, library):
# For list libraries, we don't return the fully qualified lib name. eg. arctic_skhare.test -> skhare.test
if library.startswith(self.DB_PREFIX + '_'):
return library[len(self.DB_PREFIX) + 1:]
return library
@mongo_retry
def initialize_library(self, library, lib_type=VERSION_STORE, **kwargs):
"""
Create an Arctic Library or a particular type.
Parameters
----------
library : `str`
The name of the library. e.g. 'library' or 'user.library'
lib_type : `str`
The type of the library. e.g. arctic.VERSION_STORE or arctic.TICK_STORE
Or any type registered with register_library_type
Default: arctic.VERSION_STORE
kwargs :
Arguments passed to the Library type for initialization.
"""
lib = ArcticLibraryBinding(self, library)
# check that we don't create too many namespaces
# can be disabled check_library_count=False
check_library_count = kwargs.pop('check_library_count', True)
if len(self._conn[lib.database_name].list_collection_names()) > 5000 and check_library_count:
raise ArcticException("Too many namespaces %s, not creating: %s" %
(len(self._conn[lib.database_name].list_collection_names()), library))
lib.set_library_type(lib_type)
LIBRARY_TYPES[lib_type].initialize_library(lib, **kwargs)
# Add a 10G quota just in case the user is calling this with API.
if not lib.get_quota():
lib.set_quota(10 * 1024 * 1024 * 1024)
self._cache.append('list_libraries', self._sanitize_lib_name(library))
@mongo_retry
def delete_library(self, library):
"""
Delete an Arctic Library, and all associated collections in the MongoDB.
Parameters
----------
library : `str`
The name of the library. e.g. 'library' or 'user.library'
"""
lib = ArcticLibraryBinding(self, library)
colname = lib.get_top_level_collection().name
if not [c for c in lib._db.list_collection_names(False) if re.match(r"^{}([\.].*)?$".format(colname), c)]:
logger.info('Nothing to delete. Arctic library %s does not exist.' % colname)
logger.info('Dropping collection: %s' % colname)
lib._db.drop_collection(colname)
for coll in lib._db.list_collection_names():
if coll.startswith(colname + '.'):
logger.info('Dropping collection: %s' % coll)
lib._db.drop_collection(coll)
if library in self._library_cache:
del self._library_cache[library]
del self._library_cache[lib.get_name()]
self._cache.delete_item_from_key('list_libraries', self._sanitize_lib_name(library))
def get_library(self, library):
"""
Return the library instance. Can generally use slicing to return the library:
arctic_store[library]
Parameters
----------
library : `str`
The name of the library. e.g. 'library' or 'user.library'
"""
if library in self._library_cache:
return self._library_cache[library]
try:
error = None
lib = ArcticLibraryBinding(self, library)
lib_type = lib.get_library_type()
except (OperationFailure, AutoReconnect) as e:
error = e
if error:
raise LibraryNotFoundException("Library %s was not correctly initialized in %s.\nReason: %r)" %
(library, self, error))
elif not lib_type:
raise LibraryNotFoundException("Library %s was not correctly initialized in %s." %
(library, self))
elif lib_type not in LIBRARY_TYPES:
raise LibraryNotFoundException("Couldn't load LibraryType '%s' for '%s' (has the class been registered?)" %
(lib_type, library))
instance = LIBRARY_TYPES[lib_type](lib)
self._library_cache[library] = instance
# The library official name may be different from 'library': e.g. 'library' vs 'user.library'
self._library_cache[lib.get_name()] = instance
return self._library_cache[library]
def __getitem__(self, key):
if isinstance(key, string_types):
return self.get_library(key)
else:
raise ArcticException("Unrecognised library specification - use [libraryName]")
def set_quota(self, library, quota):
"""
Set a quota (in bytes) on this user library. The quota is 'best effort',
and should be set conservatively.
Parameters
----------
library : `str`
The name of the library. e.g. 'library' or 'user.library'
quota : `int`
Advisory quota for the library - in bytes
"""
ArcticLibraryBinding(self, library).set_quota(quota)
def get_quota(self, library):
"""
Return the quota currently set on the library.
Parameters
----------
library : `str`
The name of the library. e.g. 'library' or 'user.library'
"""
return ArcticLibraryBinding(self, library).get_quota()
def check_quota(self, library):
"""
Check the quota on the library, as would be done during normal writes.
Parameters
----------
library : `str`
The name of the library. e.g. 'library' or 'user.library'
Raises
------
arctic.exceptions.QuotaExceededException if the quota has been exceeded
"""
ArcticLibraryBinding(self, library).check_quota()
def rename_library(self, from_lib, to_lib):
"""
Renames a library
Parameters
----------
from_lib: str
The name of the library to be renamed
to_lib: str
The new name of the library
"""
to_colname = to_lib
if '.' in from_lib and '.' in to_lib:
if from_lib.split('.')[0] != to_lib.split('.')[0]:
raise ValueError("Collection can only be renamed in the same database")
to_colname = to_lib.split('.')[1]
lib = ArcticLibraryBinding(self, from_lib)
colname = lib.get_top_level_collection().name
logger.info('Renaming collection: %s' % colname)
lib._db[colname].rename(to_colname)
for coll in lib._db.list_collection_names():
if coll.startswith(colname + '.'):
lib._db[coll].rename(coll.replace(colname, to_colname))
if from_lib in self._library_cache:
del self._library_cache[from_lib]
del self._library_cache[lib.get_name()]
self._cache.update_item_for_key(
'list_libraries', self._sanitize_lib_name(from_lib), self._sanitize_lib_name(to_lib))
def get_library_type(self, lib):
"""
Returns the type of the library
Parameters
----------
lib: str
the library
"""
return ArcticLibraryBinding(self, lib).get_library_type()
class ArcticLibraryBinding(object):
"""
The ArcticLibraryBinding type holds the binding between the library name and the
concrete implementation of the library.
Also provides access to additional metadata about the library
- Access to the library's top-level collection
- Enforces quota on the library
- Access to custom metadata about the library
"""
DB_PREFIX = Arctic.DB_PREFIX
TYPE_FIELD = "TYPE"
QUOTA = 'QUOTA'
quota = None
quota_countdown = 0
@classmethod
def _parse_db_lib(cls, library):
"""
Returns the canonical (database_name, library) for the passed in
string 'library'.
"""
database_name = library.split('.', 2)
if len(database_name) == 2:
library = database_name[1]
if database_name[0].startswith(cls.DB_PREFIX):
database_name = database_name[0]
else:
database_name = cls.DB_PREFIX + '_' + database_name[0]
else:
database_name = cls.DB_PREFIX
return database_name, library
def __init__(self, arctic, library):
self.arctic = arctic
self._curr_conn = self.arctic._conn
self._lock = threading.RLock()
database_name, library = self._parse_db_lib(library)
self.library = library
self.database_name = database_name
self._auth(self.arctic._conn[self.database_name])
@property
def _db(self):
with self._lock:
arctic_conn = self.arctic._conn
if arctic_conn is not self._curr_conn:
self._auth(arctic_conn[self.database_name]) # trigger re-authentication if Arctic has been reset
self._curr_conn = arctic_conn
return self.arctic._conn[self.database_name]
@property
def _library_coll(self):
return self._db[self.library]
def __str__(self):
return """<ArcticLibrary at %s, %s.%s>
%s""" % (hex(id(self)), self._db.name, self._library_coll.name, indent(str(self.arctic), 4))
def __repr__(self):
return str(self)
def __getstate__(self):
return {'arctic': self.arctic, 'library': '.'.join([self.database_name, self.library])}
def __setstate__(self, state):
return ArcticLibraryBinding.__init__(self, state['arctic'], state['library'])
@mongo_retry
def _auth(self, database):
# Get .mongopass details here
if not hasattr(self.arctic, 'mongo_host'):
return
auth = get_auth(self.arctic.mongo_host, self.arctic._application_name, database.name)
if auth:
authenticate(database, auth.user, auth.password)
def reset_auth(self):
logger.debug("reset_auth() %s" % self)
self._auth(self._db)
def get_name(self):
return self._db.name + '.' + self._library_coll.name
def get_top_level_collection(self):
"""
Return the top-level collection for the Library. This collection is to be used
for storing data.
Note we expect (and callers require) this collection to have default read-preference: primary
The read path may choose to reduce this if secondary reads are allowed.
"""
return self._library_coll
def set_quota(self, quota_bytes):
"""
Set a quota (in bytes) on this user library. The quota is 'best effort',
and should be set conservatively.
A quota of 0 is 'unlimited'
"""
self.set_library_metadata(ArcticLibraryBinding.QUOTA, quota_bytes)
self.quota = quota_bytes
self.quota_countdown = 0
def get_quota(self):
"""
Get the current quota on this user library.
"""
return self.get_library_metadata(ArcticLibraryBinding.QUOTA)
def check_quota(self):
"""
Check whether the user is within quota. Should be called before
every write. Will raise() if the library has exceeded its allotted
quota.
"""
# Don't check on every write, that would be slow
if self.quota_countdown > 0:
self.quota_countdown -= 1
return
# Re-cache the quota after the countdown
self.quota = self.get_library_metadata(ArcticLibraryBinding.QUOTA)
if self.quota is None or self.quota == 0:
self.quota = 0
return
# Figure out whether the user has exceeded their quota
library = self.arctic[self.get_name()]
stats = library.stats()
def to_gigabytes(bytes_):
return bytes_ / 1024. / 1024. / 1024.
# Have we exceeded our quota?
size = stats['totals']['size']
count = stats['totals']['count']
if size >= self.quota:
raise QuotaExceededException("Mongo Quota Exceeded: %s %.3f / %.0f GB used" % (
'.'.join([self.database_name, self.library]),
to_gigabytes(size),
to_gigabytes(self.quota)))
# Quota not exceeded, print an informational message and return
try:
avg_size = size // count if count > 1 else 100 * 1024
remaining = self.quota - size
remaining_count = remaining / avg_size
if remaining_count < 100 or float(remaining) / self.quota < 0.1:
logger.warning("Mongo Quota: %s %.3f / %.0f GB used" % (
'.'.join([self.database_name, self.library]),
to_gigabytes(size),
to_gigabytes(self.quota)))
else:
logger.info("Mongo Quota: %s %.3f / %.0f GB used" % (
'.'.join([self.database_name, self.library]),
to_gigabytes(size),
to_gigabytes(self.quota)))
# Set-up a timer to prevent us for checking for a few writes.
# This will check every average half-life
self.quota_countdown = int(max(remaining_count // 2, 1))
except Exception as e:
logger.warning("Encountered an exception while calculating quota statistics: %s" % str(e))
def get_library_type(self):
return self.get_library_metadata(ArcticLibraryBinding.TYPE_FIELD)
def set_library_type(self, lib_type):
self.set_library_metadata(ArcticLibraryBinding.TYPE_FIELD, lib_type)
@mongo_retry
def get_library_metadata(self, field):
lib_metadata = self._library_coll[self.arctic.METADATA_COLL].find_one({"_id": self.arctic.METADATA_DOC_ID})
if lib_metadata is not None:
return lib_metadata.get(field)
else:
return None
@mongo_retry
def set_library_metadata(self, field, value):
self._library_coll[self.arctic.METADATA_COLL].update_one({'_id': self.arctic.METADATA_DOC_ID},
{'$set': {field: value}}, upsert=True)
