<commit_before>from celery.task.base import Task, PeriodicTask
from celery.execute import send_task

from pymongo.objectid import ObjectId

from oyster.conf import settings
from oyster.client import get_configured_client


class UpdateTask(Task):
    """ task that updates documents """
    # results go straight to database
    ignore_result = True

    def __init__(self):
        # one client per process
        self.client = get_configured_client()


    def run(self, doc_id):
        doc = self.client.db.tracked.find_one({'_id': doc_id})
        self.client.update(doc)
        for hook in doc.get('post_update_hooks', []):
            send_task(hook, (doc_id,))
        self.client.db.status.update({}, {'$inc': {'update_queue': -1}})


class UpdateTaskScheduler(PeriodicTask):
    """ recurring task that populates the update queue """

    # 60s tick
    run_every = 60
    client = get_configured_client()

    def run(self):
        # if the update queue isn't empty, wait to add more
        # (currently the only way we avoid duplicates)
        # alternate option would be to set a _queued flag on documents
        if self.client.db.status.find_one()['update_queue']:
            return

        next_set = self.client.get_update_queue()
        for doc in next_set:
            UpdateTask.delay(doc['_id'])
            self.client.db.status.update({}, {'$inc': {'update_queue': 1}})


class ExternalStoreTask(Task):
    """ base class for tasks that push documents to an external store

        when overiding be sure to define
            external_store
                short string describing external store (eg. 's3')
            upload_document(self, doc_id, filedata, metadata)
                function that uploads the document and returns a unique ID
    """

    # results go straight to database
    ignore_result = True
    # used as a base class
    abstract = True

    def __init__(self):
        # one client per process
        self.client = get_configured_client()

    def run(self, doc_id, extract_text=lambda x: x):
        # get the document
        doc = self.client.db.tracked.find_one({'_id': ObjectId(doc_id)})
        filedata = self.client.get_version(doc['url']).read()
        text = extract_text(filedata)

        # put the document into the data store
        result = self.upload_document(doc_id, text, doc['metadata'])

        doc[self.external_store + '_id'] = result
        self.client.db.tracked.save(doc, safe=True)


    def upload_document(self, doc_id, filedata, metadata):
        """ abstract method, override on implementations """
        pass


class S3StoreTask(ExternalStoreTask):<commit_msg>pass metadata to extract_text<commit_after>from celery.task.base import Task, PeriodicTask
from celery.execute import send_task

from pymongo.objectid import ObjectId

from oyster.conf import settings
from oyster.client import get_configured_client


class UpdateTask(Task):
    """ task that updates documents """
    # results go straight to database
    ignore_result = True

    def __init__(self):
        # one client per process
        self.client = get_configured_client()


    def run(self, doc_id):
        doc = self.client.db.tracked.find_one({'_id': doc_id})
        self.client.update(doc)
        for hook in doc.get('post_update_hooks', []):
            send_task(hook, (doc_id,))
        self.client.db.status.update({}, {'$inc': {'update_queue': -1}})


class UpdateTaskScheduler(PeriodicTask):
    """ recurring task that populates the update queue """

    # 60s tick
    run_every = 60
    client = get_configured_client()

    def run(self):
        # if the update queue isn't empty, wait to add more
        # (currently the only way we avoid duplicates)
        # alternate option would be to set a _queued flag on documents
        if self.client.db.status.find_one()['update_queue']:
            return

        next_set = self.client.get_update_queue()
        for doc in next_set:
            UpdateTask.delay(doc['_id'])
            self.client.db.status.update({}, {'$inc': {'update_queue': 1}})


class ExternalStoreTask(Task):
    """ base class for tasks that push documents to an external store

        when overiding be sure to define
            external_store
                short string describing external store (eg. 's3')
            upload_document(self, doc_id, filedata, metadata)
                function that uploads the document and returns a unique ID
    """

    # results go straight to database
    ignore_result = True
    # used as a base class
    abstract = True

    def __init__(self):
        # one client per process
        self.client = get_configured_client()

    def run(self, doc_id, extract_text=lambda x: x):
        # get the document
        doc = self.client.db.tracked.find_one({'_id': ObjectId(doc_id)})
        filedata = self.client.get_version(doc['url']).read()
        text = extract_text(filedata, doc['metadata'])

        # put the document into the data store
        result = self.upload_document(doc_id, text, doc['metadata'])

        doc[self.external_store + '_id'] = result
        self.client.db.tracked.save(doc, safe=True)


    def upload_document(self, doc_id, filedata, metadata):
        """ abstract method, override on implementations """
        pass


class S3StoreTask(ExternalStoreTask):
{'fraction_non_alphanumeric': 0.09923378807699496, 'fraction_numerical': 0.0029900953092879836, 'mean_word_length': 2.7610681658468024, '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': '1251862', 'n_tokens_mistral': 1556, 'n_tokens_neox': 1485, 'n_words': 455}