<commit_before>
''' DataStore class for WorkBench '''

import pymongo
import gridfs
import hashlib
import datetime
import bson


class DataStore():

    def __init__(self, uri='mongodb://localhost/workbench', database='workbench', worker_cap=0, samples_cap=0):

        self.sample_collection = 'samples'
        self.worker_cap = worker_cap
        self.samples_cap = samples_cap

        # Get connection to mongo
        self.uri = uri
        self.db_name = database
        self.c = pymongo.MongoClient('mongodb://'+self.uri+'/'+self.db_name)
        self.db = self.c.get_default_database()

        # Get the gridfs handle
        self.gridfs_handle = gridfs.GridFS(self.db)

        # Run the periodic operations
        self.periodic_ops()

        print 'WorkBench DataStore connected: %s:%s' % (self.uri, self.db_name)

    def get_uri(self):
        return self.uri

    def store_sample(self, filename, sample_bytes, type_tag):

        sample_info = {}

        # Compute the MD5 hash
        sample_info['md5'] = hashlib.md5(sample_bytes).hexdigest()

        # Check if sample already exists
        if self.have_sample(sample_info['md5']):
            print 'Sample %s: already exists in DataStore' % (sample_info['md5'])
            return sample_info['md5']

        # Print info
        print 'Sample Storage: %.2f out of %.2f MB' % (self.sample_storage_size(), self.samples_cap)

        # Check if we need to run periodic operations
        if (self.sample_storage_size() > self.samples_cap):
            self.periodic_ops()

        # Check if we need to expire anything
        self.expire_data()

        # Okay start populating the sample for adding to the data store
        # Filename, length, import time and type_tag
        sample_info['filename'] = filename
        sample_info['length'] = len(sample_bytes)
        sample_info['import_time'] = datetime.datetime.utcnow()
        sample_info['type_tag'] = type_tag

        # Random customer for now
        import random
        sample_info['customer'] = random.choice(['Mega Corp','Huge Inc','BearTron','<NAME>'])

        # Push the file into the MongoDB GridFS
        print 'Storing Sample into Mongo GridFS'
        sample_info['__grid_fs'] = self.gridfs_handle.put(sample_bytes)
        self.db[self.sample_collection].insert(sample_info)

        # Return the sample md5
        return sample_info['md5']

    def sample_storage_size(self):
        ''' Get the storage size of the samples storage collection '''
        coll_stats = self.db.command('collStats', 'fs.chunks')
        sample_storage_size = coll_stats['size']/1024.0/1024.0
        return sample_storage_size

    def expire_data(self):
        ''' Expire data within the samples collection '''

        # Do we need to start deleting stuff?
        while (self.sample_storage_size() > self.samples_cap):

            # This should return the 'oldest' record in samples
            record = self.db[self.sample_collection].find().sort('import_time',pymongo.ASCENDING).limit(1)[0]

            # Delete it
            print 'Deleting sample: %s (%.2f MB)...' % (record['md5'], record['length']/1024.0/1024.0)
            self.db[self.sample_collection].remove({'md5': record['md5']})
            self.gridfs_handle.delete(record['__grid_fs'])

            # Print info
            print 'Sample Storage: %.2f out of %.2f MB' % (self.sample_storage_size(), self.samples_cap)

    def clean_for_serialization(self, data):
        if isinstance(data, dict):
            for k in data.keys():
                if (k.startswith('__')): del data[k]
                elif isinstance(data[k], bson.objectid.ObjectId): del data[k]
                elif isinstance(data[k], datetime.datetime):
                    data[k] = data[k].isoformat()+'Z'
                elif isinstance(data[k], dict):
                    data[k] = self.clean_for_serialization(data[k])
                elif isinstance(data[k], list):
                    data[k] = [self.clean_for_serialization(item) for item in data[k]]
        return data

    def clean_for_storage(self, data):
        data = self.data_to_unicode(data)
        if isinstance(data, dict):
            for k in data.keys():
                if k == '_id':
                    del data[k]
                    continue
                if '.' in k:
                    new_k = k.replace('.', '_')
                    data[new_k] = data[k]
                    del data[k]
                    k = new_k
                if isinstance(data[k], dict):
                    data[k] = self.clean_for_storage(data[k])
                elif isinstance(data[k], list):
                    data[k] = [self.clean_for_storage(item) for item in data[k]]
        return data

    def get_sample(self, md5):
        ''' Get the sample from the data store '''
        sample_info = self.db[self.sample_collection].find_one({'md5':md5})
        if (not sample_info):
            raise RuntimeError('Sample not found: %s ' % (md5))

        # Get the raw bytes from GridFS (note: this could fail)
        try:
            grid_fs_id = sample_info['__grid_fs']
            sample_info = self.clean_for_serialization(sample_info)
            sample_info.update({'raw_bytes':self.gridfs_handle.get(grid_fs_id).read()})
            return sample_info
        except gridfs.errors.CorruptGridFile:
            # If we don't have the gridfs files, delete the entry from samples
            self.db[self.sample_collection].update({'md5':md5}, {'md5':None})
            raise RuntimeError('Sample not found: %s ' % (md5))

    def have_sample(self, md5):
        ''' See if the data store has this sample '''

        # The easiest thing is to simply get the sample and if that
        # succeeds than return True, else return False
        try:
            self.get_sample(md5)
            return True
        except:
            return False

    def store_work_results(self, results, collection, md5):
        results['md5'] = md5
        results['__time_stamp'] = datetime.datetime.utcnow()

        # Fixme: Occasionally a capped collection will not let you update with a 
        #        larger object, if you have MongoDB 2.6 or above this shouldn't
        #        really happen, so for now just kinda punting and giving a message.
        try:
            self.db[collection].update({'md5':md5}, self.clean_for_storage(results), True)
        except pymongo.errors.OperationFailure:
            print 'Not updating exising object in capped collection...(upgrade to mongodb 2.6)'

    def get_work_results(self, collection, md5):
        return self.db[collection].find_one({'md5':md5})

    def all_sample_md5s(self, type_tag=None):
        if type_tag:
            cursor = self.db[self.sample_collection].find({'type_tag':type_tag},{'md5':1, '_id':0})
        else:
            cursor = self.db[self.sample_collection].find({},{'md5':1, '_id':0})
        return [ match.values()[0] for match in cursor ]

    def periodic_ops(self):
        ''' Run periodic operations on the the data store
            Things like making sure collections are capped
            and indexes are set up '''

        # Get all the collections in the workbench database
        all_c = self.db.collection_names()

        # Remove collections that we don't want to cap
        try:
            all_c.remove('system.indexes')
            all_c.remove('fs.chunks')
            all_c.remove('fs.files')
            all_c.remove(self.sample_collection)
        except ValueError:
            pass

        # Convert collections to capped if desired
        if (self.worker_cap):
            size = self.worker_cap * pow(1024, 2)  # MegaBytes per collection
            for collection in all_c:
                self.db.command('convertToCapped', collection, size=size)

        # Loop through all collections ensuring they have an index on MD5s
        for collection in all_c:
            self.db[collection].ensure_index('md5')

        # Add required index for samples collection
        self.db[self.sample_collection].create_index('import_time')

    # Helper functions
    def to_unicode(self, s):

        # Fixme: This is total horseshit
        if isinstance(s, unicode):
            return s
        if isinstance(s, str):
            return unicode(s, errors='ignore')

        # Just return the original object
        return s

    def data_to_unicode(self, data):
        if isinstance(data, dict):
            return {self.to_unicode(k): self.to_unicode(v) for k, v in data.iteritems()}
        if isinstance(data, list):
            return [self.to_unicode(l) for l in data]
        else:
            return self.to_unicode(data)
<commit_msg>fixing periodic ops to only run every 5 min<commit_after>
''' DataStore class for WorkBench '''

import pymongo
import gridfs
import hashlib
import datetime
import bson
import time


class DataStore():

    def __init__(self, uri='mongodb://localhost/workbench', database='workbench', worker_cap=0, samples_cap=0):

        self.sample_collection = 'samples'
        self.worker_cap = worker_cap
        self.samples_cap = samples_cap

        # Get connection to mongo
        self.uri = uri
        self.db_name = database
        self.c = pymongo.MongoClient('mongodb://'+self.uri+'/'+self.db_name)
        self.db = self.c.get_default_database()

        # Get the gridfs handle
        self.gridfs_handle = gridfs.GridFS(self.db)

        # Run the periodic operations
        self.last_ops_run = time.time()
        self.periodic_ops()

        print 'WorkBench DataStore connected: %s:%s' % (self.uri, self.db_name)

    def get_uri(self):
        return self.uri

    def store_sample(self, filename, sample_bytes, type_tag):

        sample_info = {}

        # Compute the MD5 hash
        sample_info['md5'] = hashlib.md5(sample_bytes).hexdigest()

        # Check if sample already exists
        if self.have_sample(sample_info['md5']):
            print 'Sample %s: already exists in DataStore' % (sample_info['md5'])
            return sample_info['md5']

        # Run the periodic operations
        self.periodic_ops()

        # Check if we need to expire anything
        self.expire_data()

        # Okay start populating the sample for adding to the data store
        # Filename, length, import time and type_tag
        sample_info['filename'] = filename
        sample_info['length'] = len(sample_bytes)
        sample_info['import_time'] = datetime.datetime.utcnow()
        sample_info['type_tag'] = type_tag

        # Random customer for now
        import random
        sample_info['customer'] = random.choice(['Mega Corp','Huge Inc','BearTron','Dorseys Mom'])

        # Push the file into the MongoDB GridFS
        sample_info['__grid_fs'] = self.gridfs_handle.put(sample_bytes)
        self.db[self.sample_collection].insert(sample_info)

        # Print info
        print 'Sample Storage: %.2f out of %.2f MB' % (self.sample_storage_size(), self.samples_cap)

        # Return the sample md5
        return sample_info['md5']

    def sample_storage_size(self):
        ''' Get the storage size of the samples storage collection '''
        coll_stats = self.db.command('collStats', 'fs.chunks')
        sample_storage_size = coll_stats['size']/1024.0/1024.0
        return sample_storage_size

    def expire_data(self):
        ''' Expire data within the samples collection '''

        # Do we need to start deleting stuff?
        while (self.sample_storage_size() > self.samples_cap):

            # This should return the 'oldest' record in samples
            record = self.db[self.sample_collection].find().sort('import_time',pymongo.ASCENDING).limit(1)[0]

            # Delete it
            print 'Deleting sample: %s (%.2f MB)...' % (record['md5'], record['length']/1024.0/1024.0)
            self.db[self.sample_collection].remove({'md5': record['md5']})
            self.gridfs_handle.delete(record['__grid_fs'])

            # Print info
            print 'Sample Storage: %.2f out of %.2f MB' % (self.sample_storage_size(), self.samples_cap)

    def clean_for_serialization(self, data):
        if isinstance(data, dict):
            for k in data.keys():
                if (k.startswith('__')): del data[k]
                elif isinstance(data[k], bson.objectid.ObjectId): del data[k]
                elif isinstance(data[k], datetime.datetime):
                    data[k] = data[k].isoformat()+'Z'
                elif isinstance(data[k], dict):
                    data[k] = self.clean_for_serialization(data[k])
                elif isinstance(data[k], list):
                    data[k] = [self.clean_for_serialization(item) for item in data[k]]
        return data

    def clean_for_storage(self, data):
        data = self.data_to_unicode(data)
        if isinstance(data, dict):
            for k in data.keys():
                if k == '_id':
                    del data[k]
                    continue
                if '.' in k:
                    new_k = k.replace('.', '_')
                    data[new_k] = data[k]
                    del data[k]
                    k = new_k
                if isinstance(data[k], dict):
                    data[k] = self.clean_for_storage(data[k])
                elif isinstance(data[k], list):
                    data[k] = [self.clean_for_storage(item) for item in data[k]]
        return data

    def get_sample(self, md5):
        ''' Get the sample from the data store '''
        sample_info = self.db[self.sample_collection].find_one({'md5':md5})
        if (not sample_info):
            raise RuntimeError('Sample not found: %s ' % (md5))

        # Get the raw bytes from GridFS (note: this could fail)
        try:
            grid_fs_id = sample_info['__grid_fs']
            sample_info = self.clean_for_serialization(sample_info)
            sample_info.update({'raw_bytes':self.gridfs_handle.get(grid_fs_id).read()})
            return sample_info
        except gridfs.errors.CorruptGridFile:
            # If we don't have the gridfs files, delete the entry from samples
            self.db[self.sample_collection].update({'md5':md5}, {'md5':None})
            raise RuntimeError('Sample not found: %s ' % (md5))

    def have_sample(self, md5):
        ''' See if the data store has this sample '''

        # The easiest thing is to simply get the sample and if that
        # succeeds than return True, else return False
        try:
            self.get_sample(md5)
            return True
        except:
            return False

    def store_work_results(self, results, collection, md5):
        results['md5'] = md5
        results['__time_stamp'] = datetime.datetime.utcnow()

        # Fixme: Occasionally a capped collection will not let you update with a 
        #        larger object, if you have MongoDB 2.6 or above this shouldn't
        #        really happen, so for now just kinda punting and giving a message.
        try:
            self.db[collection].update({'md5':md5}, self.clean_for_storage(results), True)
        except pymongo.errors.OperationFailure:
            print 'Not updating exising object in capped collection...(upgrade to mongodb 2.6)'

    def get_work_results(self, collection, md5):
        return self.db[collection].find_one({'md5':md5})

    def all_sample_md5s(self, type_tag=None):
        if type_tag:
            cursor = self.db[self.sample_collection].find({'type_tag':type_tag},{'md5':1, '_id':0})
        else:
            cursor = self.db[self.sample_collection].find({},{'md5':1, '_id':0})
        return [ match.values()[0] for match in cursor ]

    def periodic_ops(self):
        ''' Run periodic operations on the the data store
            Things like making sure collections are capped
            and indexes are set up '''

        # Only run every 5 minutes
        if (time.time() - self.last_ops_run) < 300:
            return

        # Reset last ops run
        self.last_ops_run = time.time()
        print 'Running Periodic Ops'

        # Get all the collections in the workbench database
        all_c = self.db.collection_names()

        # Remove collections that we don't want to cap
        try:
            all_c.remove('system.indexes')
            all_c.remove('fs.chunks')
            all_c.remove('fs.files')
            all_c.remove(self.sample_collection)
        except ValueError:
            pass

        # Convert collections to capped if desired
        if (self.worker_cap):
            size = self.worker_cap * pow(1024, 2)  # MegaBytes per collection
            for collection in all_c:
                self.db.command('convertToCapped', collection, size=size)

        # Loop through all collections ensuring they have an index on MD5s
        for collection in all_c:
            self.db[collection].ensure_index('md5')

        # Add required index for samples collection
        self.db[self.sample_collection].create_index('import_time')

    # Helper functions
    def to_unicode(self, s):

        # Fixme: This is total horseshit
        if isinstance(s, unicode):
            return s
        if isinstance(s, str):
            return unicode(s, errors='ignore')

        # Just return the original object
        return s

    def data_to_unicode(self, data):
        if isinstance(data, dict):
            return {self.to_unicode(k): self.to_unicode(v) for k, v in data.iteritems()}
        if isinstance(data, list):
            return [self.to_unicode(l) for l in data]
        else:
            return self.to_unicode(data)

{'fraction_non_alphanumeric': 0.10627276867030966, 'fraction_numerical': 0.009050546448087431, 'mean_word_length': 2.2183550100751055, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 7, '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': '6303567', 'n_tokens_mistral': 5165, 'n_tokens_neox': 4920, 'n_words': 1354}