<commit_before># Copyright (c) 2016, the ElectrumX authors
#
# All rights reserved.
#
# See the file "LICENCE" for information about the copyright
# and warranty status of this software.

'''Backend database abstraction.

The abstraction needs to be improved to not heavily penalise LMDB.
'''

import os
from functools import partial

from lib.util import subclasses, increment_byte_string


def open_db(name, db_engine, for_sync):
    '''Returns a database handle.'''
    for db_class in subclasses(Storage):
        if db_class.__name__.lower() == db_engine.lower():
            db_class.import_module()
            return db_class(name, for_sync)

    raise RuntimeError('unrecognised DB engine "{}"'.format(db_engine))


class Storage(object):
    '''Abstract base class of the DB backend abstraction.'''

    def __init__(self, name, for_sync):
        self.is_new = not os.path.exists(name)
        self.open(name, create=self.is_new, for_sync=for_sync)

    @classmethod
    def import_module(cls):
        '''Import the DB engine module.'''
        raise NotImplementedError

    def open(self, name, create):
        '''Open an existing database or create a new one.'''
        raise NotImplementedError

    def close(self):
        '''Close an existing database.'''
        raise NotImplementedError

    def get(self, key):
        raise NotImplementedError

    def put(self, key, value):
        raise NotImplementedError

    def write_batch(self):
        '''Return a context manager that provides `put` and `delete`.

        Changes should only be committed when the context manager
        closes without an exception.
        '''
        raise NotImplementedError

    def iterator(self, prefix=b'', reverse=False):
        '''Return an iterator that yields (key, value) pairs from the
        database sorted by key.

        If `prefix` is set, only keys starting with `prefix` will be
        included.  If `reverse` is True the items are returned in
        reverse order.
        '''
        raise NotImplementedError


class LevelDB(Storage):
    '''LevelDB database engine.'''

    @classmethod
    def import_module(cls):
        import plyvel
        cls.module = plyvel

    def open(self, name, create, for_sync):
        mof = 1024 if for_sync else 256
        self.db = self.module.DB(name, create_if_missing=create,
                                 max_open_files=mof, compression=None)
        self.close = self.db.close
        self.get = self.db.get
        self.put = self.db.put
        self.iterator = self.db.iterator
        self.write_batch = partial(self.db.write_batch, transaction=True)


class RocksDB(Storage):
    '''RocksDB database engine.'''

    @classmethod
    def import_module(cls):
        import rocksdb
        cls.module = rocksdb

    def open(self, name, create, for_sync):
        mof = 1024 if for_sync else 256
        compression = "no"
        compression = getattr(self.module.CompressionType,
                              compression + "_compression")
        options = self.module.Options(create_if_missing=create,
                                      compression=compression,
                                      target_file_size_base=33554432,
                                      max_open_files=mof)
        self.db = self.module.DB(name, options)
        self.get = self.db.get
        self.put = self.db.put

    def close(self):
        # PyRocksDB doesn't provide a close method; hopefully this is enough
        self.db = None
        import gc
        gc.collect()

    class WriteBatch(object):
        def __init__(self, db):
            self.batch = RocksDB.module.WriteBatch()
            self.db = db

        def __enter__(self):
            return self.batch

        def __exit__(self, exc_type, exc_val, exc_tb):
            if not exc_val:
                self.db.write(self.batch)

    def write_batch(self):
        return RocksDB.WriteBatch(self.db)

    class Iterator(object):
        def __init__(self, db, prefix, reverse):
            self.it = db.iteritems()
            self.reverse = reverse
            self.prefix = prefix
            # Whether we are at the first item
            self.first = True

        def __iter__(self):
            prefix = self.prefix
            if self.reverse:
                prefix = increment_byte_string(prefix)
                self.it = reversed(self.it)<commit_msg>Remove all RocksDB references on close<commit_after># Copyright (c) 2016, the ElectrumX authors
#
# All rights reserved.
#
# See the file "LICENCE" for information about the copyright
# and warranty status of this software.

'''Backend database abstraction.

The abstraction needs to be improved to not heavily penalise LMDB.
'''

import os
from functools import partial

from lib.util import subclasses, increment_byte_string


def open_db(name, db_engine, for_sync):
    '''Returns a database handle.'''
    for db_class in subclasses(Storage):
        if db_class.__name__.lower() == db_engine.lower():
            db_class.import_module()
            return db_class(name, for_sync)

    raise RuntimeError('unrecognised DB engine "{}"'.format(db_engine))


class Storage(object):
    '''Abstract base class of the DB backend abstraction.'''

    def __init__(self, name, for_sync):
        self.is_new = not os.path.exists(name)
        self.open(name, create=self.is_new, for_sync=for_sync)

    @classmethod
    def import_module(cls):
        '''Import the DB engine module.'''
        raise NotImplementedError

    def open(self, name, create):
        '''Open an existing database or create a new one.'''
        raise NotImplementedError

    def close(self):
        '''Close an existing database.'''
        raise NotImplementedError

    def get(self, key):
        raise NotImplementedError

    def put(self, key, value):
        raise NotImplementedError

    def write_batch(self):
        '''Return a context manager that provides `put` and `delete`.

        Changes should only be committed when the context manager
        closes without an exception.
        '''
        raise NotImplementedError

    def iterator(self, prefix=b'', reverse=False):
        '''Return an iterator that yields (key, value) pairs from the
        database sorted by key.

        If `prefix` is set, only keys starting with `prefix` will be
        included.  If `reverse` is True the items are returned in
        reverse order.
        '''
        raise NotImplementedError


class LevelDB(Storage):
    '''LevelDB database engine.'''

    @classmethod
    def import_module(cls):
        import plyvel
        cls.module = plyvel

    def open(self, name, create, for_sync):
        mof = 1024 if for_sync else 256
        self.db = self.module.DB(name, create_if_missing=create,
                                 max_open_files=mof, compression=None)
        self.close = self.db.close
        self.get = self.db.get
        self.put = self.db.put
        self.iterator = self.db.iterator
        self.write_batch = partial(self.db.write_batch, transaction=True)


class RocksDB(Storage):
    '''RocksDB database engine.'''

    @classmethod
    def import_module(cls):
        import rocksdb
        cls.module = rocksdb

    def open(self, name, create, for_sync):
        mof = 1024 if for_sync else 256
        compression = "no"
        compression = getattr(self.module.CompressionType,
                              compression + "_compression")
        options = self.module.Options(create_if_missing=create,
                                      compression=compression,
                                      target_file_size_base=33554432,
                                      max_open_files=mof)
        self.db = self.module.DB(name, options)
        self.get = self.db.get
        self.put = self.db.put

    def close(self):
        # PyRocksDB doesn't provide a close method; hopefully this is enough
        self.db = self.get = self.put = None
        import gc
        gc.collect()

    class WriteBatch(object):
        def __init__(self, db):
            self.batch = RocksDB.module.WriteBatch()
            self.db = db

        def __enter__(self):
            return self.batch

        def __exit__(self, exc_type, exc_val, exc_tb):
            if not exc_val:
                self.db.write(self.batch)

    def write_batch(self):
        return RocksDB.WriteBatch(self.db)

    class Iterator(object):
        def __init__(self, db, prefix, reverse):
            self.it = db.iteritems()
            self.reverse = reverse
            self.prefix = prefix
            # Whether we are at the first item
            self.first = True

        def __iter__(self):
            prefix = self.prefix
            if self.reverse:
                prefix = increment_byte_string(prefix)
                self.it = reversed(self.it)
{'fraction_non_alphanumeric': 0.08800631840234684, 'fraction_numerical': 0.005867087893489789, 'mean_word_length': 2.1410347271438694, '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': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4466331', 'n_tokens_mistral': 2512, 'n_tokens_neox': 2390, 'n_words': 730}