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()