diff --git a/r2/r2/models/promo_metrics.py b/r2/r2/models/promo_metrics.py index 5976a3e05..c3126f2fd 100644 --- a/r2/r2/models/promo_metrics.py +++ b/r2/r2/models/promo_metrics.py @@ -20,6 +20,8 @@ # Inc. All Rights Reserved. ############################################################################### +from itertools import product + from r2.lib.db import tdb_cassandra from r2.lib.utils import tup from r2.models.subreddit import DefaultSR @@ -66,3 +68,54 @@ class PromoMetrics(tdb_cassandra.View): values_by_sr[fp] = values_by_sr.get(fp, 0) + values_by_sr[''] del(values_by_sr['']) cls._set_values(metric_name, values_by_sr) + + +class LocationPromoMetrics(tdb_cassandra.View): + _use_db = True + _write_consistency_level = tdb_cassandra.CL.QUORUM + _read_consistency_level = tdb_cassandra.CL.ONE + _extra_schema_creation_args = { + "default_validation_class": tdb_cassandra.IntegerType(), + } + + @classmethod + def _rowkey(cls, location): + fields = [location.country, location.region, location.metro] + return '-'.join(map(lambda field: field or '', fields)) + + @classmethod + def _column_name(cls, sr): + return sr.name + + @classmethod + def get(cls, srs, locations): + srs, srs_is_single = tup(srs, ret_is_single=True) + locations, locations_is_single = tup(locations, ret_is_single=True) + is_single = srs_is_single and locations_is_single + + rowkeys = {location: cls._rowkey(location) for location in locations} + columns = {sr: cls._column_name(sr) for sr in srs} + rcl = cls._read_consistency_level + metrics = cls._cf.multiget(rowkeys.values(), columns.values(), + read_consistency_level=rcl) + ret = {} + + for sr, location in product(srs, locations): + rowkey = rowkeys[location] + column = columns[sr] + impressions = metrics.get(rowkey, {}).get(column, 0) + ret[(sr, location)] = impressions + + if is_single: + return ret.values()[0] + else: + return ret + + @classmethod + def set(cls, metrics): + wcl = cls._write_consistency_level + with cls._cf.batch(write_consistency_level=wcl) as b: + for location, sr, impressions in metrics: + rowkey = cls._rowkey(location) + column = {cls._column_name(sr): impressions} + b.insert(rowkey, column)