<commit_before>import pymongo

from motor.motor_tornado import MotorClient
from pymongo import MongoClient

from .conf import config
from .exceptions import DbError


_db = None


def connect_to_db(async=True):
    db_conf = config['db']

    with_auth = db_conf.get('user') and db_conf.get('password')
    if with_auth:
        auth = '{user}:{password}@'.format(**db_conf)
        if db_conf.get('authDbName'):
            auth_db = '/' + db_conf['authDbName']
    else:
        auth = ''
        auth_db = ''

    dest = '{host}:{port}'.format(**db_conf)

    conn_str = 'mongodb://' + auth + dest + auth_db
    client_cls = MotorClient if async else MongoClient
    db_client = client_cls(conn_str)

    db = getattr(db_client, db_conf['dbName'])

    return db_client, db


class ModelField:
    def __init__(self):
        pass


class StringField(ModelField):
    def __init__(self):
        pass


class IntegerField(ModelField):
    def __init__(self):
        pass


class ModelMeta(type):
    def __new__(cls, name, bases, attrs):
        if 'collection_name' not in attrs:
            attrs['collection_name'] = name.lower()

        return super().__new__(cls, name, bases, attrs)


class Model(metaclass=ModelMeta):
    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            setattr(self, k, v)

    @classmethod
    def find(cls, **kwargs):
        return Collection(cls, cls._get_db()).find(**kwargs)

    @classmethod
    def find_one(cls, **kwargs):
        return Collection(cls, cls._get_db()).find_one(**kwargs)

    @classmethod
    def _get_db(cls):
        if hasattr(cls, 'db'):
            return cls.db
        raise Exception()


def get_db(conn_str=None, db_name=None):
    global _db

    if _db is None:
        _db = MongoClient(conn_str)[db_name]

    return _db


class Collection:
    def __init__(self, model_cls, db):
        if not issubclass(model_cls, Model):
            raise ValueError(
                'Class "%s" is not a subclass of "logster.db.Model" class',
                model_cls.__name__)

        self.model_cls = model_cls
        self.collection = db[model_cls.collection_name]
        self.result = None<commit_msg>Improve database retrieval in collection<commit_after>import pymongo

from motor.motor_tornado import MotorClient
from pymongo import MongoClient

from .conf import config
from .exceptions import DbError


_db = None


def connect_to_db(async=True):
    db_conf = config['db']

    with_auth = db_conf.get('user') and db_conf.get('password')
    if with_auth:
        auth = '{user}:{password}@'.format(**db_conf)
        if db_conf.get('authDbName'):
            auth_db = '/' + db_conf['authDbName']
    else:
        auth = ''
        auth_db = ''

    dest = '{host}:{port}'.format(**db_conf)

    conn_str = 'mongodb://' + auth + dest + auth_db
    client_cls = MotorClient if async else MongoClient
    db_client = client_cls(conn_str)

    db = getattr(db_client, db_conf['dbName'])

    return db_client, db


class ModelField:
    def __init__(self):
        pass


class StringField(ModelField):
    def __init__(self):
        pass


class IntegerField(ModelField):
    def __init__(self):
        pass


class ModelMeta(type):
    def __new__(cls, name, bases, attrs):
        if 'collection_name' not in attrs:
            attrs['collection_name'] = name.lower()

        return super().__new__(cls, name, bases, attrs)


class Model(metaclass=ModelMeta):
    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            setattr(self, k, v)

    @classmethod
    def find(cls, **kwargs):
        return Collection(cls).find(**kwargs)

    @classmethod
    def find_one(cls, **kwargs):
        return Collection(cls).find_one(**kwargs)

    # @classmethod
    # def _get_db(cls):
    #     if hasattr(cls, 'db'):
    #         return cls.db
    #     return None


def get_db(conn_str=None, db_name=None):
    global _db

    if _db is None:
        _db = MongoClient(conn_str)[db_name]

    return _db


class Collection:
    def __init__(self, model_cls, db=None):
        if not issubclass(model_cls, Model):
            raise ValueError(
                'Class "%s" is not a subclass of "logster.db.Model" class',
                model_cls.__name__)

        self.model_cls = model_cls
        self.collection = (db or get_db())[model_cls.collection_name]
        self.result = None
{'fraction_non_alphanumeric': 0.11564625850340136, 'fraction_numerical': 0.0, 'mean_word_length': 2.6697171381031612, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '709612', 'n_tokens_mistral': 1449, 'n_tokens_neox': 1392, 'n_words': 347}