Redis in action, chapter 1, StackOverflow mini functionality

redis-in-action-chap1.py

#!/usr/bin/python2.6

import redis
import time
import random

ONE_WEEK_IN_SECONDS = 7 * 86400
VOTE_SCORE = 432
ARTICLES_PER_PAGE = 25

def article_vote(conn, user, article, upvote=True):
	cutoff = time.time() - ONE_WEEK_IN_SECONDS

	if conn.zscore('time:', article) < cutoff:
		return

	article_id = article.split(':')[-1]

	if upvote:
		key = 'voted:' + article_id
		opponent_key = 'downvoted:' + article_id

		article_vote_key = 'votes'
		article_opponent_vote_key = 'downvotes'
		vote_score = VOTE_SCORE
	else:
		key = 'downvoted:' + article_id
		opponent_key = 'voted:' + article_id

		article_vote_key = 'downvotes'
		article_opponent_vote_key = 'votes'
		vote_score = VOTE_SCORE * -1

	if conn.sismember(key, user):
		print "DEBUG: %s already %s for %s" % (user, key, article)
		return

	if conn.sismember(opponent_key, user):
		print "DEBUG: %s switched their mind: %s -> %s for %s" % (user, opponent_key, key, article)
		conn.smove(opponent_key, key, user)
		conn.zincrby('score:', article, vote_score * 2)	# *2 for switch vote
		conn.hincrby(article, article_opponent_vote_key, -1)
		conn.hincrby(article, article_vote_key, 1)
	else:
		conn.sadd(key, user)
		conn.zincrby('score:', article, vote_score)
		conn.hincrby(article, article_vote_key, 1)

def post_article(conn, user, title, link):
	article_id = str(conn.incr('article:'))

	voted = 'voted:' + article_id
	conn.sadd(voted, user)
	conn.expire(voted, ONE_WEEK_IN_SECONDS + 1)

	now = time.time()
	article = 'article:' + article_id
	conn.hmset(article, {
		'title': title,
		'link': link,
		'poster': user,
		'time': now,
		'votes': 1,
		'downvotes': 0,
	})

	conn.zadd('score:', article, now + VOTE_SCORE)
	conn.zadd('time:', article, now)

	return article_id

def get_articles(conn, page, order='score:'):
	start = max(page-1, 0) * ARTICLES_PER_PAGE
	end = start + ARTICLES_PER_PAGE - 1

	ids = conn.zrevrange(order, start, end)
	articles = []
	for id in ids:
		article_data = conn.hgetall(id)
		article_data['id'] = id
		articles.append(article_data)

	return articles

def add_remove_groups(conn, article_id, to_add=[], to_remove=[]):
	article = 'article:' + article_id

	for group in to_add:
		conn.sadd('group:' + group, article)

	for group in to_remove:
		conn.srem('group:' + group, article)

def get_group_articles(conn, group, page, order='score:'):
	key = order + group
	if not conn.exists(key):
		conn.zinterscore(key, ['group:' + group, order], aggregate='max')
		conn.expire(key, 60)
	return get_articles(conn, page, key)

def print_article_status(article_id):
	print "*" * 40
	print conn.hgetall('article:' + article_id)
	print conn.zscore('score:', 'article:' + article_id)
	print conn.zscore('time:', 'article:' + article_id)
	print conn.smembers('voted:' + article_id)
	print conn.smembers('downvoted:' + article_id)
	print "*" * 40

def do_random_vote(conn, user, start_article_id, end_article_id):
	# select article
	article_id = random.randint(start_article_id, end_article_id)

	# upvote / downvote
	if random.randint(0, 1) == 1:
		upvote = True
	else:
		upvote = False

	article_vote(conn, user, 'article:%d' % article_id, upvote=upvote)

	return (article_id, upvote)

def apply_random_vote(conn, start_article_id, end_article_id):
	user_list = ['user:%d' % x for x in range(100)]

	# first votes
	for user in user_list:
		do_random_vote(conn, user, start_article_id, end_article_id)
		
	# second votes
	for user in user_list:
		do_random_vote(conn, user, start_article_id, end_article_id)

def test_post_article(conn):
	article_id = post_article(conn, 'user:83271', 'Go To Statement??', 'http://goo.gl/kZUSu')

	print_article_status(article_id)

	return article_id

def test_vote(conn, article_id):
	article_key = 'article:' + article_id

	# with new voted user 12345, upvote
	article_vote(conn, 'user:12345', article_key)

	print_article_status(article_id)

	# with already upvoted user
	article_vote(conn, 'user:12345', article_key)

	print_article_status(article_id)

	# with new voted user 54321, downvote
	article_vote(conn, 'user:54321', article_key, upvote=False)

	print_article_status(article_id)

	# user 12345 switched vote, upvote -> downvote
	article_vote(conn, 'user:12345', article_key, upvote=False)

	print_article_status(article_id)

	# user 54321 switched vote, downvote -> upvote
	article_vote(conn, 'user:54321', article_key)

	print_article_status(article_id)

if __name__ == "__main__":
	conn = redis.Redis('search-crw-dev2')

	# article_id = test_post_article(conn)
	# test_vote(conn, article_id)

	conn.flushdb()
	for i in range(10):
		article_id = post_article(conn, 'user:83271', 'Go To Statement??', 'http://goo.gl/kZUSu')
		#article_id = test_post_article(conn)

	#article_id => 1 ~ 10
	apply_random_vote(conn, 1, 10)

	page_articles_by_score = get_articles(conn, 1, order='score:')
	page_articles_by_time = get_articles(conn, 1, order='time:')

	print "order by score:"
	print page_articles_by_score

	print "order by time:"
	print page_articles_by_time

	votes_list = map(lambda x: int(x['votes']), page_articles_by_score)
	downvotes_list = map(lambda x: int(x['downvotes']), page_articles_by_score)
	votes_count = sum(votes_list)
	downvotes_count = sum(downvotes_list)

	# 10 posts (with poster vote), 100 users * 2 votes => 210 votes
	# second votes(100 users) can duplicated or switched, and these votes are not counted to whole votes
	# so, (votes + downvotes) must be 210 - (duplicate + switched)
	print "%s votes, %s downvotes, total %s" % (votes_count, downvotes_count, votes_count + downvotes_count)