def user_logged_in(shard, user_id):
"""Logs in a user to a shard. Always returns the current user ID."""
login_record = None
if user_id:
# Re-login the user if they somehow lost their browser state and
# needed to reload the page. This assumes the cookie was okay.
login_record = models.LoginRecord.get_by_id(user_id)
if login_record and not login_record.online:
def txn():
login_record = models.LoginRecord.get_by_id(user_id)
assert login_record
login_record.online = True
login_record.put()
logging.debug('Re-logging-in user_id=%r to shard=%r',
login_record.user_id, shard)
ndb.transaction(txn)
# User is logging in for the first time or somehow state was deleted.
if not login_record:
login_record = models.LoginRecord(
key=ndb.Key(models.LoginRecord._get_kind(), models.human_uuid()),
shard_id=shard,
online=True)
login_record.put()
logging.debug('Logged-in new user_id=%r to shard=%r',
login_record.user_id, shard)
invalidate_user_cache(shard)
return login_record.user_id
def post(self):
address = self.get_required('address', str, '')
secret = self.get_required('secret', str, '')
action = self.get_required('action', str).lower()
opt_out = action == 'opt-out'
opt_in = action == 'opt-in'
def txn():
if not address:
raise ndb.Rollback()
email_record = models.EmailRecord.get_by_id(address)
if not email_record or email_record.secret != secret:
raise ndb.Rollback()
if opt_out:
logging.debug('Opting out of email. address=%r', address)
email_record.global_opt_out = True
elif opt_in:
logging.debug('Opting-in to email. address=%r', address)
email_record.global_opt_out = False
else:
raise ndb.Rollback()
email_record.put()
ndb.transaction(txn)
self.redirect('/email?address=%s&secret=%s' % (address, secret))
def test_count_dec(self):
comment = Comment.query(Comment.article == model.Key('Article', 1)).get()
#1/0
ndb.transaction(comment.key.delete, xg=True)
self.submit_deferred()
get_count = Aggregation.query(Aggregation.refkey == model.Key('Article', 1),
Aggregation.field == 'article_comment_count').get().value
self.assertEqual(self.count-1, get_count)
def test_sum_sub(self):
comment = Comment.query(Comment.article == model.Key('Article', 1)).get()
last_rating = comment.rating
ndb.transaction(comment.key.delete, xg=True)
self.submit_deferred()
get_total_sum = Aggregation.query(Aggregation.refkey == model.Key('Article', 1),
Aggregation.field == 'article_rating_sum').get().value
self.assertEqual(self.total_rating-last_rating, get_total_sum)
def test_existing_sum_change(self):
comment = Comment.query(Comment.article == model.Key('Article', 1)).get()
last_rating = comment.rating
comment.rating = 7.5
ndb.transaction(comment.put, xg=True)
self.submit_deferred()
get_total_sum = Aggregation.query(Aggregation.refkey == model.Key('Article', 1),
Aggregation.field == 'article_rating_sum').get().value
get_total_count = Aggregation.query(Aggregation.refkey == model.Key('Article', 1),
Aggregation.field == 'article_comment_count').get().value
self.assertEqual(self.total_rating - last_rating + comment.rating, get_total_sum)
self.assertEqual(self.count, get_total_count)
def user_logged_out(shard, user_id):
"""Notifies other users that the given user has logged out of a shard."""
def txn():
login_record = models.LoginRecord.get_by_id(user_id)
if not login_record:
raise ndb.Rollback()
if not login_record.online:
raise ndb.Rollback()
login_record.online = False
login_record.put()
return login_record
login_record = ndb.transaction(txn)
if not login_record:
logging.warning('Tried to log out user_id=%r from shard=%r, '
'but LoginRecord did not exist', user_id, shard)
return
posts.insert_post(
shard,
archive_type=models.Post.USER_LOGOUT,
nickname=login_record.nickname,
user_id=user_id,
body='%s has left' % login_record.nickname)
invalidate_user_cache(shard)
logging.debug('Logged out user_id=%r from shard=%r', user_id, shard)
def handle_request(self):
shard_id = None
while True:
def txn():
shard_id = models.human_uuid()
shard = models.Shard.get_by_id(
shard_id, use_cache=False, use_memcache=False)
if shard:
raise ndb.Rollback()
shard = models.Shard(id=shard_id)
shard.put()
return shard_id
shard_id = ndb.transaction(txn)
if shard_id:
break
# For dev_appserver use the relative path. Otherwise use a sub-domain.
if config.is_dev_appserver:
self.redirect('/chat/' + shard_id)
else:
host_parts = self.request.host.split('.', 1)
target_url = '%s://%s.%s/' % (
self.request.scheme, shard_id, host_parts[1])
self.redirect(target_url)
def rollback(self, migration, force=False):
def callback(steps, migration, force):
for step in reversed(steps):
step.rollback(migration, force=force)
try:
if self.fake_transaction:
callback(self.steps, migration, force)
else:
ndb.transaction(lambda:callback(self.steps, migration, force), xg=True)
#except datastore_errors.TransactionFailedError:
except Exception:
if force or self.ignore_errors in ('rollback', 'all'):
logging.exception("Ignored error in transaction while rolling back")
return
migration.fail()
raise
def apply(self, migration, force=False):
def callback(steps, migration, force):
for step in steps:
step.apply(migration, force=force)
try:
if self.fake_transaction:
callback(self.steps, migration, force)
else:
ndb.transaction(lambda:callback(self.steps, migration, force), xg=True)
#except datastore_errors.TransactionFailedError:
except Exception, err:
if force or self.ignore_errors in ('apply', 'all'):
logging.exception("Ignored error in transaction while applying")
return
migration.fail()
raise
def setUp(self):
super(AggregationTestCase, self).setUp()
# Initialising TestComments objects
for i in range(1, 3):
Article(content = rand_text()).put()
for i in range(1, 15):
comment = Comment(article = Article.get_by_id(random.randrange(1,3)).key,
body= rand_text(),
rating= random.randrange(0,11)
)
ndb.transaction(comment.put, xg=True)
self.submit_deferred()
comment = Comment.query(Comment.article == model.Key('Article', 1))
self.count = comment.count()
self.total_rating = 0
for field in comment:
self.total_rating += field.rating
def update_read_state(topic_dict, user_id):
"""Updates a user's read state for a given set of topics.
Will inherit transaction state on user_id's LoginRecord if this is
called from within an existing transition.
Args:
topic_dict: Maps topic shard IDs to the new sequence number to assign
for that shard.
user_id: User ID that is being updated.
"""
# TODO(bslatkin): Consider validating the topic list provided here
# to ensure they are actually associated with the current logged-in shard.
# The possible damage here is restricted to the user so we don't care much.
read_state_keys = [
ndb.Key(models.LoginRecord._get_kind(), user_id,
models.ReadState._get_kind(), topic)
for topic in topic_dict]
def txn():
read_state_list = ndb.get_multi(read_state_keys)
to_put = []
for key, read_state in zip(read_state_keys, read_state_list):
next_read_sequence = topic_dict[key.id()]
if read_state is None:
read_state = models.ReadState(key=key)
read_state.last_read_sequence = next_read_sequence
else:
read_state.last_read_sequence = max(
read_state.last_read_sequence, next_read_sequence)
to_put.append(read_state)
ndb.put_multi(to_put)
if ndb.in_transaction():
txn()
else:
ndb.transaction(txn)
def test_aggretation_tran(self):
def _put(self, **ctx_options):
request = webapp2.Request.blank("/")
request.method = 'GET'
request.app = test_app.app
webapp2._local.request = request
super(self.__class__, self)._put(**ctx_options)
if self.rating == 3:
import time; time.sleep(1)
get_total_sum = Aggregation.query(Aggregation.refkey == model.Key('Article', 1),
Aggregation.field == 'article_rating_sum').get().value
get_total_count = Aggregation.query(Aggregation.refkey == model.Key('Article', 1),
Aggregation.field == 'article_comment_count').get().value
old_put = Comment.put
Comment.put = _put
comment = Comment.query(Comment.article == model.Key('Article', 1)).get()
comment.rating = 3
ndb.transaction(comment.put, xg=True)
self.submit_deferred()
self.assertEqual(self.total_rating, get_total_sum)
self.assertEqual(self.count, get_total_count)
Comment.put = old_put
def change_presence(shard, user_id, nickname, accepted_terms,
sounds_enabled, retrying, email_address):
"""Changes the presence for a user."""
def txn():
last_nickname = None
user_connected = True
login = models.LoginRecord.get_by_id(user_id)
if not login:
login = models.LoginRecord(
key=ndb.Key(models.LoginRecord._get_kind(), user_id),
shard_id=shard)
elif only_active_users(login):
# This is a heartbeat presence check
user_connected = False
if maybe_update_token(login, force=retrying):
logging.debug(
'Issuing channel token: user_id=%r, shard=%r, force=%r',
user_id, shard, retrying)
if nickname:
# This is a potential nickname change. Right now the client
# always sends the nickname on every request, so we need to
# check for the difference to detect a rename.
last_nickname = login.nickname
login.nickname = nickname
if accepted_terms:
# This is a ToS acceptance
login.accepted_terms_version = config.terms_version
login.online = True
login.sounds_enabled = sounds_enabled
login.email_address = email_address or None
login.put()
return last_nickname, user_connected, login.browser_token
last_nickname, user_connected, browser_token = ndb.transaction(txn)
# Invalidate the cache so the nickname will be updated next time
# someone requests the roster.
invalidate_user_cache(shard)
message = None
archive_type = None
if nickname and last_nickname and last_nickname != nickname:
message = '%s has changed their nickname to %s' % (
last_nickname, nickname)
archive_type = models.Post.USER_UPDATE
logging.debug('User update user_id=%r, shard=%r', user_id, shard)
elif user_connected:
message = '%s has joined' % nickname
archive_type = models.Post.USER_LOGIN
logging.debug('User joined: user_id=%r, shard=%r', user_id, shard)
else:
logging.debug('User heartbeat: user_id=%r to shard=%r',
user_id, shard)
if archive_type:
posts.insert_post(
shard,
archive_type=archive_type,
nickname=nickname,
user_id=user_id,
body=message)
else:
# As long as users are heart-beating, we should be running a
# cleanup task for this shard.
enqueue_cleanup_task(shard)
return user_connected, browser_token
def send_digest_email(email_address, sequence_number):
"""Sends a digest email to the user with only what's changed."""
email_record = models.EmailRecord.get_or_insert(
email_address,
secret=models.human_uuid())
if email_record.global_opt_out:
logging.debug('Not sending email to globally opted out address=%r',
email_address)
return
if email_record.sequence_number > sequence_number:
logging.warning(
'Saw email digest task for address=%r, sequence_number=%r but '
'email record already at sequence_number=%r; dropping task',
email_address, sequence_number, email_record.sequence_number)
return
# Find all shards the user participates in with this email address.
query = models.LoginRecord.query()
query = query.filter(models.LoginRecord.email_address == email_address)
login_record_list = query.fetch(1000)
shard_set = set(u.shard_id for u in login_record_list)
# Generate the template rendering params for each shard and all of its
# topics based on the email's read state for each topic. Do this in
# parallel since none of them are interdependent.
futures_dict = {}
for root_shard_id in shard_set:
futures_dict[root_shard_id] = get_topic_info(
root_shard_id, email_address)
ndb.Future.wait_all(futures_dict.values())
shard_list = []
for root_shard_id, future in futures_dict.iteritems():
topic_info_dict = future.get_result()
if not topic_info_dict.get('topic_list'):
logging.debug('No updates for shard=%r, email=%r',
root_shard_id, email_address)
continue
shard_list.append(topic_info_dict)
# Mark the email as processed so we can send future digests.
def txn():
email_record = models.EmailRecord.get_by_id(email_address)
email_record.sequence_number = sequence_number + 1
email_record.previous_notified_time = \
email_record.last_notified_time
email_record.last_notified_time = datetime.datetime.now()
email_record.put()
return email_record
email_record = ndb.transaction(txn)
if not shard_list:
logging.debug('No topics to digest for %s', email_address)
return
total_topics = 0
total_updates = 0
for shard_dict in shard_list:
total_topics += shard_dict.get('total_topics', 0)
total_updates += shard_dict.get('total_updates', 0)
subject = 'What\'s new:'
if total_topics == 1:
subject += ' 1 topic'
elif total_topics > 1:
subject += ' %d topics' % total_topics
if total_topics and total_updates:
subject += ','
if total_updates == 1:
subject += ' 1 update'
elif total_updates > 1:
subject += ' %d updates' % total_updates
# Render the email content and send the email. We do this last so if
# there are any bugs in the rendering code or sending step that it does
# not result in us repeatedly sending emails to users. The transaction
# above acts as a guard on this task.
context = dict(
cache_buster=config.version_id,
email_record=email_record,
email_resource_host_prefix=config.email_resource_host_prefix,
email_title=subject,
shard_list=shard_list,
site_name=config.site_name
)
text_data = template.render('templates/digest_email.txt', context)
html_data = template.render('templates/digest_email_output.html', context)
sender = config.notification_from_email
logging.debug('Sending email digest to=%r, sender=%r, subject=%r',
email_address, sender, subject)
message = mail.EmailMessage(
sender=sender,
to=email_address,
subject=subject,
body=text_data,
html=html_data)
message.send()
def apply_posts(shard=None,
insertion_post_id=None,
lease_seconds=10,
max_tasks=20):
"""Applies a set of pending posts to a shard.
If shard is None then this function will apply mods for whatever is the
first shard it can find in the pull task queue.
insertion_post_id is the post_id that first caused this apply task to be
enqueued. This task will retry until it applies the insertion_post_id
itself or it can confirm that the insertion_post_id has already been
applied. insertion_post_id may be empty if the apply task is not associated
with a particular post (such as cronjobs/cleanup tasks).
"""
# Do not use caching for NDB in this task queue worker.
ctx = ndb.get_context()
ctx.set_cache_policy(lambda x: False)
ctx.set_memcache_policy(lambda x: False)
# Fetch the new Posts to put in sequence.
queue = taskqueue.Queue(config.pending_queue)
# When no shard is specified, process the first tag we find.
task_list = []
if not shard:
task_list.extend(queue.lease_tasks(lease_seconds, 1))
if not task_list:
logging.debug('apply_posts with no specific shard found no tasks')
return
params = task_list[0].extract_params()
shard = params['shard']
logging.debug('apply_posts with no specific shard found shard=%r',
shard)
# Clear the dirty bit on this shard to start the time horizon.
dirty_bit(shard, clear=True)
# Find tasks pending for the current shard.
task_list.extend(
queue.lease_tasks_by_tag(lease_seconds, max_tasks, tag=str(shard)))
receipt_key_list = []
new_topic = None
for task in task_list:
params = task.extract_params()
# Extract the new topic shard associated with this task, if any. The
# last one wins. If all of the found posts have already been applied,
# then topic assignment will be ignored.
new_topic = params.get('new_topic') or new_topic
post_id_list = params.get('post_ids')
if post_id_list is None:
# This may happen on replica shards if it turns out there are no
# unapplied post IDs but an apply task still ran.
post_id_list = []
elif not isinstance(post_id_list, list):
post_id_list = [post_id_list]
for post_id in post_id_list:
receipt_key = ndb.Key(
models.Post._get_kind(), post_id,
models.Receipt._get_kind(), shard)
receipt_key_list.append(receipt_key)
receipt_list = ndb.get_multi(receipt_key_list)
# Some tasks may be in the pull queue that were already put in sequence.
# So ignore these and only apply the new ones.
unapplied_receipts = [
models.Receipt(key=k)
for k, r in zip(receipt_key_list, receipt_list)
if r is None]
unapplied_post_ids = [r.post_id for r in unapplied_receipts]
# Double check if we think there should be work to apply but we didn't find
# any. This will force the apply task to retry immediately if the post task
# was not found. This can happen when the pull queue's consistency is
# behind.
if not unapplied_receipts and insertion_post_id:
receipt_key = ndb.Key(
models.Post._get_kind(), insertion_post_id,
models.Receipt._get_kind(), shard)
receipt = receipt_key.get()
if receipt:
logging.warning(
'No post application to do for shard=%r, but post_id=%r '
'already applied; doing nothing in this task',
shard, insertion_post_id)
new_topic = None
# Do not 'return' here. We need to increment the shard sequence or
# else tasks will not run for this shard in the future because of
# de-duping.
else:
raise base.Error('No post application to do for shard=%r, but'
'post_id=%r has not been applied; will retry' %
(shard, insertion_post_id))
now = datetime.datetime.now()
def txn():
shard_record = models.Shard.get_by_id(shard)
# TODO(bslatkin): Just drop this task entirely if the shard cannot
# be found. Could happen for old shards that were cleaned up.
assert shard_record
# One of the tasks in this batch has a topic assignment. Apply it here.
if new_topic:
logging.debug('Changing topic from %r to %r',
shard_record.current_topic, new_topic)
shard_record.current_topic = new_topic
shard_record.topic_change_time = now
new_sequence_numbers = list(xrange(
shard_record.sequence_number,
shard_record.sequence_number + len(unapplied_receipts)))
shard_record.sequence_number += max(1, len(unapplied_receipts))
# Write post references that point at the newly sequenced posts.
to_put = [shard_record]
for receipt, sequence in zip(unapplied_receipts, new_sequence_numbers):
to_put.append(models.PostReference(
id=sequence,
parent=shard_record.key,
post_id=receipt.post_id))
# Update the receipt entity here; it will be written outside this
# transaction, since these receipts may span multiple entity
# groups.
receipt.sequence = sequence
# Enqueue replica posts transactionally, to make sure everything
# definitely will get copied over to the replica shard.
if shard_record.current_topic:
enqueue_post_task(shard_record.current_topic, unapplied_post_ids)
ndb.put_multi(to_put)
return shard_record, new_sequence_numbers
# Have this only attempt a transaction a single time. If the transaction
# fails the task queue will retry this task within 4 seconds. Because
# apply tasks are always named by the current Shard.sequence_number we
# can be reasonably sure that no other apply task for this shard will be
# running concurrently when this fails.
shard_record, new_sequence_numbers = ndb.transaction(txn, retries=1)
replica_shard = shard_record.current_topic
logging.debug('Applied %d posts for shard=%r, sequence_numbers=%r',
len(unapplied_receipts), shard, new_sequence_numbers)
futures = []
# Save receipts for all the posts.
futures.extend(ndb.put_multi_async(unapplied_receipts))
# Notify all logged in users of the new posts.
futures.append(notify_posts(
shard, unapplied_post_ids, sequence_numbers=new_sequence_numbers))
# Replicate posts to a topic shard.
if replica_shard:
logging.debug('Replicating source shard=%r to replica shard=%r',
shard, replica_shard)
futures.append(enqueue_apply_task(replica_shard))
# Success! Delete the tasks from this queue.
queue.delete_tasks(task_list)
# Always run one more apply task to clean up any posts that came in
# while this transaction was processing.
if dirty_bit(shard, check=True):
futures.append(enqueue_apply_task(shard))
# Wait on all pending futures in case they raise errors.
ndb.Future.wait_all(futures)
# For root shards, add shard cleanup task to check for user presence and
# cause notification of user logouts if the channel API did not detect the
# user closing the connection.
if not shard_record.root_shard:
presence.enqueue_cleanup_task(shard)
def handle_request(self, shard_id):
shard = models.Shard.get_by_id(shard_id)
if not shard:
# If the shard doesn't exist, then just create it. Makes it
# ridiculously easy for people to create a new chat with the name
# of their choice.
def txn():
shard = models.Shard.get_by_id(
shard_id, use_cache=False, use_memcache=False)
if not shard:
shard = models.Shard(id=shard_id)
shard.put()
return shard
shard = ndb.transaction(txn)
# Do not allow users to directly access topic shards.
if shard.root_shard:
# TODO(bslatkin): Make this pretty.
self.response.set_status(404)
return
nickname = 'Anonymous'
email_address = ''
first_login = True
must_accept_terms = True
sounds_enabled = True
if 'shards' in self.session:
# TODO(bslatkin): Reuse presence code here.
user_id = self.session['shards'].get(shard_id)
if user_id:
login_record = models.LoginRecord.get_by_id(user_id)
if login_record and login_record.shard_id == shard_id:
email_address = login_record.email_address
nickname = login_record.nickname
first_login = False
must_accept_terms = bool(
login_record.accepted_terms_version !=
config.terms_version)
sounds_enabled = login_record.sounds_enabled
else:
# Pre-populate the user's settings info from the newest other
# shard the user has.
all_user_ids = self.session['shards'].values()
all_login_records = ndb.get_multi([
ndb.Key(models.LoginRecord._get_kind(), user_id)
for user_id in all_user_ids])
all_login_records = [r for r in all_login_records if r]
all_login_records.sort(key=lambda x: x.last_update_time)
if all_login_records:
login_record = all_login_records[-1]
email_address = login_record.email_address
nickname = login_record.nickname
sounds_enabled = login_record.sounds_enabled
context = {
'email_address': email_address or '',
'first_login': first_login,
'must_accept_terms': must_accept_terms,
'nickname': xml.sax.saxutils.unescape(nickname),
'page_name': 'chatroom',
'shard_id': shard_id,
'short_url': self.request.path_url,
'sounds_enabled': sounds_enabled,
'xsrf_token': self.session['xsrf_token'],
}
context['params'] = json.dumps(context)
self.render('chatroom.html', context)
