From 3c031c8db6e3f10a60ab0fa471a5555f6b68321f Mon Sep 17 00:00:00 2001 From: Brian Simpson Date: Sun, 12 Jun 2016 09:14:22 -0700 Subject: [PATCH] TryLater.multi_handle: Document context manager magic --- r2/r2/models/trylater.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/r2/r2/models/trylater.py b/r2/r2/models/trylater.py index 4cf665aff7..6c54989dd3 100644 --- a/r2/r2/models/trylater.py +++ b/r2/r2/models/trylater.py @@ -75,23 +75,24 @@ class TryLater(tdb_cassandra.View): _write_consistency_level = tdb_cassandra.CL.QUORUM _compare_with = TIME_UUID_TYPE - @classmethod - def multi_ready(cls, rowkeys, cutoff=None): - if cutoff is None: - cutoff = datetime.datetime.utcnow() - return cls._cf.multiget(rowkeys, - column_finish=cutoff, - column_count=tdb_cassandra.max_column_count) - @classmethod @contextlib.contextmanager def multi_handle(cls, rowkeys, cutoff=None): if cutoff is None: cutoff = datetime.datetime.utcnow() - ready = cls.multi_ready(rowkeys, cutoff) + + ready = cls._cf.multiget( + rowkeys, + column_finish=cutoff, + column_count=tdb_cassandra.max_column_count, + ) + + # return the rows to the context caller yield ready + + # on context __exit__ cleanup all the ready columns for system, items in ready.iteritems(): - cls._remove(system, items.keys()) + cls._cf.remove(system, items.keys()) @classmethod def search(cls, rowkey, when):