Saving comments and submissions from Reddit using Elasticsearch

Question

I wrote a script to save comments and submissions from Reddit using Elasticsearch. Is there any way to improve the speed of inserting new items into Elasticsearch?

import praw
import sys
from elasticsearch import Elasticsearch

subreddit_queue_file = "subreddit_queue.txt"

def pop_task():
    task = ""

    with open(subreddit_queue_file, "r") as f:
        lines = f.read().split()

    if len(lines) == 0:
        task = None
    else:
        task = lines.pop(0)

    lines = "\n".join(lines) + '\n'

    with open(subreddit_queue_file, "w") as f:
        f.write(lines)

    return task

def push_task(task):
    with open(subreddit_queue_file, "a") as f:
        f.write(task + '\n')


fetch_limit = 100

r = praw.Reddit('Comment Scraper 1.0 by u/kerseykyle')


def save_submission(es, submission):
    content = {
              "id": str(submission.id),
              "author": str(submission.author),
              "domain": str(submission.domain),
              "is_self": submission.is_self,
              "num_comments": int(submission.num_comments),
              "over_18": submission.over_18,
              "permalink": submission.permalink,
              "score": int(submission.score),
              "selftext": submission.selftext,
              "selftext_html": submission.selftext_html,
              "subreddit": str(submission.subreddit),
              "subreddit_id": str(submission.subreddit_id),
              "thumbnail": str(submission.thumbnail),
              "title": submission.title,
              "url": submission.url,
              "edited": str(submission.edited),
              "distinguished": str(submission.distinguished),
              "stickied": submission.stickied
            }

    res = es.search(index="_all", body={"query": {"term": {"id": str(submission.id)}}})
    if res["hits"]["total"] == 0:
        es.index(index=str(submission.subreddit).lower(), doc_type="submission", id=str(submission.id), body=content)


def save_comment(es, comment, submission_id):
    content={
             "id": str(comment.id),
             "author": str(comment.author),
             "name": str(comment.name),
             "body": comment.body,
             "body_html": comment.body_html,
             "created_utc": comment.created_utc,
             "created": comment.created,
             "edited": int(comment.edited),
             "link_id": str(comment.link_id),
             "parent_id": str(comment.parent_id),
             "score": comment.score,
             "subreddit": str(comment.subreddit),
             "subreddit_id": str(comment.subreddit_id),
             "submission_id": str(submission_id)
             }

    res = es.search(index="_all", body={"query": {"term": {"id": str(comment.id)}}})
    if res["hits"]["total"] == 0:
        es.index(index=str(comment.subreddit).lower(), doc_type="comment", id=str(comment.id), body=content)

def get_reddit_submissions(task):

    good_count = 1.0
    bad_count = 1.0
    submission_count = 0
    is_link = False

    es = Elasticsearch()

    if task.startswith("http://www.reddit.com/r/"):
        submissions = [r.get_submission(task)]
        is_link = True
    else:
        subreddit = r.get_subreddit(task)
        submissions = subreddit.get_hot(limit=fetch_limit)
    try:
        for submission in submissions:

            save_submission(es, submission)

            if is_link:
                submission.replace_more_comments(limit=None, threshold=0)

            flat_comments = praw.helpers.flatten_tree(submission.comments)
            submission_count += 1
            print "[%d %s %s]"%(submission_count, submission.subreddit, submission.id), (submission.title).encode("UTF-8")
            for comment in flat_comments:
                try:
                    save_comment(es, comment, submission.id)
                    good_count += 1
                except AttributeError:
                    bad_count += 1
                    pass


    except KeyboardInterrupt:
        print "comments insterted [%%%f]"%((bad_count/good_count)*100), good_count
        sys.exit(1)

    print "comments insterted [%%%f]"%((bad_count/good_count)*100), good_count

task = sys.argv[1]
push_task(task)

task = pop_task()

while task != None:
    if task == "":
        continue

    get_reddit_submissions(task)

    task = pop_task()

Answer 1

The main downfalls of speed in your code are going to come from constantly writing and reading a file.

I think the best option to reduce this and ultimately speed up your program would be to

1. Read the file (as you do).
2. Split the file (as you do).
3. Store the array globally.
4. Do not write back to the file.

This is going to speed up your code because you aren't reading and writing the file every single time you call pop_task; you only are reading once in the beginning.

Following the steps above, instead of getting the next task from the file in pop_task, you would simply pop the next file off the array of lines.

Here is pop_task rewritten:

def pop_task():
    task = ""

    if len(tasks) == 0:
        task = None
    else:
        task = tasks.pop(0)

    return task

Note: I switched the name lines to tasks as tasks makes more sense in a global scope.

For doing the only file read at the beginning of your code, you just simply need to load global tasks with the tasks of the file subreddit_queue_file before the infinite loop.

---

At the bottom of your code, you should be encasing your code in

if __name__ == "__main__":
    [code]

Read this StackOverflow post for why you should do this.

---

These two lines:

fetch_limit = 100

r = praw.Reddit('Comment Scraper 1.0 by u/kerseykyle')

Should not be floating in the middle of your code. They should either be up at the top where you keep your variables, or they should be down at the bottom, where the main application code is going.