Cassandra view class for roadblocks.

r2/r2/models/promo.py

@@ -25,6 +25,7 @@ from datetime import datetime, timedelta
 from uuid import uuid1
 import json
 
+from pycassa.types import CompositeType
 from pylons import g, c
 
 from r2.lib import filters
@@ -32,7 +33,7 @@ from r2.lib.cache import sgm
 from r2.lib.db import tdb_cassandra
 from r2.lib.db.thing import Thing, NotFound
 from r2.lib.memoize import memoize
-from r2.lib.utils import Enum, to_datetime
+from r2.lib.utils import Enum, to_datetime, to_date
 from r2.models.subreddit import Subreddit
 
 
@@ -247,3 +248,72 @@ class LiveAdWeights(object):
     def clear(cls, sr_id):
         """Clear ad information from the Subreddit with ID `sr_id`"""
         cls.set_all_from_weights({sr_id: []})
+
+
+class PromotedLinkRoadblock(tdb_cassandra.View):
+    _use_db = True
+    _connection_pool = 'main'
+    _read_consistency_level = tdb_cassandra.CL.ONE
+    _write_consistency_level = tdb_cassandra.CL.QUORUM
+    _compare_with = CompositeType(
+        tdb_cassandra.DateType(),
+        tdb_cassandra.DateType(),
+    )
+
+    @classmethod
+    def _column(cls, start, end):
+        start, end = map(to_datetime, [start, end])
+        return {(start, end): ''}
+
+    @classmethod
+    def _dates_from_key(cls, key):
+        start, end = map(to_date, key)
+        return start, end
+
+    @classmethod
+    def add(cls, sr, start, end):
+        rowkey = sr._id36
+        column = cls._column(start, end)
+        now = datetime.now(g.tz).date()
+        ndays = (to_date(end) - now).days + 7
+        ttl = timedelta(days=ndays).total_seconds()
+        cls._set_values(rowkey, column, ttl=ttl)
+
+    @classmethod
+    def remove(cls, sr, start, end):
+        rowkey = sr._id36
+        column = cls._column(start, end)
+        cls._remove(rowkey, column)
+
+    @classmethod
+    def is_roadblocked(cls, sr, start, end):
+        rowkey = sr._id36
+        start, end = map(to_date, [start, end])
+
+        # retrieve columns for roadblocks starting before end
+        try:
+            columns = cls._cf.get(rowkey, column_finish=(to_datetime(end),),
+                                  column_count=tdb_cassandra.max_column_count)
+        except tdb_cassandra.NotFoundException:
+            return False
+
+        for key in columns.iterkeys():
+            rb_start, rb_end = cls._dates_from_key(key)
+
+            # check for overlap, end dates not inclusive
+            if (start < rb_end) and (rb_start < end):
+                return (rb_start, rb_end)
+        return False
+
+    @classmethod
+    def get_roadblocks(cls):
+        ret = []
+        q = cls._cf.get_range()
+        rows = list(q)
+        srs = Subreddit._byID36([id36 for id36, columns in rows], data=True)
+        for id36, columns in rows:
+            sr = srs[id36]
+            for key in columns.iterkeys():
+                start, end = cls._dates_from_key(key)
+                ret.append((sr.name, start, end))
+        return ret