mongolock.py

import os
import time
import logging
from datetime import datetime

import pymongo
from mongoengine.connection import DEFAULT_CONNECTION_NAME, get_db

logger = logging.getLogger(os.path.basename(__file__))


class MongoLock:

    collection = 'locks'
    acquire_recheck = 5
    expire = 15

    def _get_collection(self):
        db = get_db(DEFAULT_CONNECTION_NAME)
        return db[self.collection]

    def __init__(self, chat_id, account_id):

        coll = self._get_collection()
        coll.create_index([('acquiredAt', pymongo.ASCENDING)],
                          expireAfterSeconds=self.expire)

        coll.create_index([('chat_id', pymongo.ASCENDING),
                           ('account_id', pymongo.ASCENDING)],
                          unique=True)

        self.chat_id = chat_id
        self.account_id = account_id

    def acquire(self):

        coll = self._get_collection()

        while True:

            try:
                coll.insert_one({
                    'chat_id': self.chat_id,
                    'account_id': self.account_id,
                    'acquiredAt': datetime.utcnow()
                })
            except pymongo.errors.DuplicateKeyError:
                pass
            else:
                logger.debug('LOCK ( %s - %s ) acquire - OK',
                             self.chat_id,
                             self.account_id)
                return True

            logger.warn('LOCK ( %s - %s ) acquire - FAIL, sleep %s',
                        self.chat_id,
                        self.account_id,
                        self.acquire_recheck)

            time.sleep(self.acquire_recheck)

    def release(self):

        coll = self._get_collection()
        coll.remove({
            'chat_id': self.chat_id,
            'account_id': self.account_id
        })

    def __enter__(self):
        self.acquire()

    def __exit__(self, *args):
        self.release()