from . import errors from .base import BaseObject from .collection import MontyCollection from .types import encode_, string_types INVALID_CHARS = ("$", "\0", "\x00") class MontyDatabase(BaseObject): def __init__(self, client, name, codec_options=None, write_concern=None): super(MontyDatabase, self).__init__( codec_options or client.codec_options, write_concern or client.write_concern ) self._client = client self._name = name self._components = (self,) def __eq__(self, other): if isinstance(other, self.__class__): return self._client == other.client and self._name == other.name return NotImplemented def __ne__(self, other): return not self == other def __repr__(self): return "MontyDatabase({!r}, {!r})".format(self._client, self._name) def __getattr__(self, name): if name.startswith("_"): raise AttributeError( "MontyDatabase has no attribute {0!r}. To access the {0}" " collection, use database[{0!r}].".format(name) ) return self.get_collection(name) def __getitem__(self, key): return self.get_collection(key) @property def name(self): """ Return database's name. """ return self._name @property def client(self): return self._client def collection_names(self): return self.client._storage.collection_list(self) list_collection_names = collection_names def create_collection(self, name, codec_options=None, write_concern=None, **kwargs): """ Create a collection or raise an error if it already exists. """ if self.client._storage.collection_exists(self, name): error_msg = "collection {} already exists".format(encode_(name)) raise errors.CollectionInvalid(error_msg) collection = self.get_collection(name, codec_options, write_concern, **kwargs) self.client._storage.collection_create(self, name) return collection def drop_collection(self, name_or_collection): if isinstance(name_or_collection, MontyCollection): name = name_or_collection.name elif not isinstance(name_or_collection, string_types): raise TypeError("name_or_collection must be an instance of basestring") else: name = name_or_collection self.client._storage.collection_drop(self, name) def get_collection(self, name, codec_options=None, write_concern=None, **kwargs): """ Return a collection if valid, otherwise throw an error. """ if not name: raise errors.OperationFailure("Collection name may not be blank") for c in INVALID_CHARS: if c in name: raise errors.OperationFailure( "Collection name contains special characters: {}".format(name) ) if name.startswith("system."): raise errors.OperationFailure( "Invalid prefix in collection name: {}".format(name) ) return MontyCollection( self, name, False, codec_options, write_concern, **kwargs )