diff --git a/heman/api/cch/mongo_curve_backend.py b/heman/api/cch/mongo_curve_backend.py index 2ef90bf..0908269 100644 --- a/heman/api/cch/mongo_curve_backend.py +++ b/heman/api/cch/mongo_curve_backend.py @@ -9,15 +9,9 @@ def __init__(self, mongodb=None): self._mongodb = mongodb or mongo.db def get_cursor_db(self, collection, query): - return self._mongodb[collection].find( - query, - fields={'_id': False, 'datetime': True, 'season': True, 'ai': True}, - ).sort([ - ('datetime', ASCENDING), - # Sorting dupped times when changing from - # summer (season=1) to winter (season=0) - ('season', DESCENDING), - ]) + return self._mongodb[collection].aggregate( + query + )['result'] def build_query( self, @@ -27,7 +21,7 @@ def build_query( **extra_filter ): - query = { + match_query = { 'name': {'$regex': '^{}'.format(cups[:20])}, # KLUDGE: datetime is naive but is stored in mongo as UTC, # if we pass dates as local, we will be comparing to the equivalent @@ -35,7 +29,25 @@ def build_query( 'datetime': {'$gte': as_naive(start), '$lt': as_naive(end)} } - query.update(extra_filter) + match_query.update(extra_filter) + + query = [ + {'$match': match_query }, + {'$group': { + '_id': {'datetime': '$datetime', 'name': '$name'}, + 'datetime': {'$first': '$datetime'}, + 'ai': {'$first': '$ai'}, + 'season': {'$first': '$season'}, + } + }, + {'$sort': { + 'datetime': ASCENDING, + # Sorting dupped times when changing from + # summer (season=1) to winter (season=0) + 'season': DESCENDING + } + } + ] return query @@ -46,7 +58,7 @@ def get_curve(self, curve_type, start, end, cups=None): for x in result: yield dict( - x, + season=x['season'], datetime=as_naive(x['datetime']), ai=float(x['ai']), ) diff --git a/tests/test_cch.py b/tests/test_cch.py index bf9bb37..e8ad46e 100644 --- a/tests/test_cch.py +++ b/tests/test_cch.py @@ -4,6 +4,7 @@ from testdata.curves import ( tg_cchfact_existing_points, tg_cchfact_NOT_existing_points_BUT_f1, + tg_cchfact_p1_repeated_records, ) from yamlns import ns @@ -160,3 +161,30 @@ def test_get_curve_p1_mongo(self, yaml_snapshot): yaml_snapshot(ns( result=list(result) )) + + def test_get_curve_p1_mongo_repeated_records(self, yaml_snapshot): + backend = MongoCurveBackend(get_mongo_instance()) + result = backend.get_curve( + curve_type=TgCchP1Repository(backend), + start=localisodate('2024-01-06'), + end=localisodate('2024-01-07'), + cups=tg_cchfact_p1_repeated_records['cups'], + ) + + yaml_snapshot(ns( + result=list(result) + )) + + def test_get_curve_p1_timescale_repeated_records(self, yaml_snapshot): + backend = TimescaleCurveBackend() + + result = backend.get_curve( + curve_type=TgCchP1Repository(backend), + start=localisodate('2024-01-06'), + end=localisodate('2024-01-07'), + cups=tg_cchfact_p1_repeated_records['cups'], + ) + + yaml_snapshot(ns( + result=[x for x in result] + )) \ No newline at end of file