mongodb_client_shared_instance.py

# -*- coding: utf-8 -*-

from __future__ import unicode_literals, print_function, absolute_import, division

from django.conf import settings
from multiprocessing import Lock
from pymongo import MongoClient
import six

__all__ = ["MongoDB"]

MONGO_DB = settings.MONGO_DB


class SingletonMetaclass(type):

    def __new__(mcs, name, bases, attrs):
        # Assume the target class is created (i.e. this method to be called) in the main thread.
        cls = super(SingletonMetaclass, mcs).__new__(mcs, name, bases, attrs)
        cls.__shared_instance_lock__ = Lock()
        # cls.__shared_instances_pool__ = {}
        return cls

    def __call__(cls, db, collection):
        with cls.__shared_instance_lock__:
            # a = {}
            # a.setdefault("a", []).append(1)
            # a.setdefault("a", []).append(2)
            # a["a"] == [1, 2]

            # if collection not in cls.__shared_instances_pool__:
            #     cls.__shared_instances_pool__[collection] = super(SingletonMetaclass, cls).__call__(db, collection)
            # return cls.__shared_instances_pool__[collection]

            # return cls.__shared_instances_pool__.setdefault(collection, super(SingletonMetaclass, cls).__call__(db, collection))

            try:
                return cls.__shared_instance__
            except AttributeError:
                cls.__shared_instance__ = super(SingletonMetaclass, cls).__call__(db, collection)
                return cls.__shared_instance__


# @six.add_metaclass(SingletonMetaclass)
class MongoDB(object):
    """ Provides a RAII wrapper for PyMongo db connections.

    Available collection functions limited to those in
    attributes_to_pass. Number of simultaneous connection
    users also tracked.

    Example:
    with MongoDB(db="iCHEF_db", collection="invoice") as m:
        m.insert(...)
        m.find(...)
        m.update(... , upsert=True)

    with MongoDB(db="iCHEF_db", collection="cancelled_invoice") as m:
        m.insert(...)
        m.find(...)
        m.update(... , upsert=True)
    """

    ATTRIBUTES_TO_PASS = ("update", "insert", "count", "find", "delete_one", "delete_many", "update_one", "update_many")
    num_users = 0
    _shared_mongo_client = None
    _shared_mongo_client_lock = Lock()
    # client = None
    # num_users_lock = Lock()

    def __init__(self, db, collection):
        self._db_name = db
        self._collection_name = collection
        self.collection = None
        # MongoDB.client = MongoDB.client or MongoClient(**MONGO_DB)

    def __enter__(self):
        with MongoDB._shared_mongo_client_lock:
            if MongoDB._shared_mongo_client is None:
                print("create client")
                MongoDB._shared_mongo_client = MongoClient(**MONGO_DB)
            MongoDB.num_users += 1
            print("add", self._collection_name, MongoDB.num_users)
        self.collection = MongoDB._shared_mongo_client[self._db_name][self._collection_name]
        return self

    def __exit__(self, *args, **kwargs):
        self.collection = None
        with MongoDB._shared_mongo_client_lock:
            MongoDB.num_users -= 1
            print("sub", self._collection_name, MongoDB.num_users)
            if MongoDB.num_users is 0:
                MongoDB._shared_mongo_client.close()
                MongoDB._shared_mongo_client = None
                print("client close")

    # def use_collection(self, collection):
    #     return MongoDBCollection(self, collection)

    def __getattr__(self, attr):
        # print(attr, self.ATTRIBUTES_TO_PASS)
        # raw_input()
        if attr in self.ATTRIBUTES_TO_PASS:
            return getattr(self.collection, attr)
        else:
            raise ValueError("")