Code improvements to time processing routines (#1010)

Author: rzats

requirements.txt

@@ -11,3 +11,4 @@ scipy==1.6.2
 tenacity==7.0.0
 newrelic
 epiweeks==2.1.2
+typing-extensions

src/server/_params.py

@@ -7,7 +7,7 @@
 
 
 from ._exceptions import ValidationFailedException
-from .utils import days_in_range, weeks_in_range, guess_time_value_is_day
+from .utils import days_in_range, weeks_in_range, guess_time_value_is_day, guess_time_value_is_week, TimeValues, days_to_ranges, weeks_to_ranges
 
 
 def _parse_common_multi_arg(key: str) -> List[Tuple[str, Union[bool, Sequence[str]]]]:
@@ -108,7 +108,7 @@ def parse_single_source_signal_arg(key: str) -> SourceSignalPair:
 @dataclass
 class TimePair:
     time_type: str
-    time_values: Union[bool, Sequence[Union[int, Tuple[int, int]]]]
+    time_values: Union[bool, TimeValues]
 
     @property
     def is_week(self) -> bool:
@@ -129,6 +129,16 @@ def count(self) -> float:
             return sum(1 if isinstance(v, int) else weeks_in_range(v) for v in self.time_values)
         return sum(1 if isinstance(v, int) else days_in_range(v) for v in self.time_values)
 
+    def to_ranges(self):
+        """
+        returns this pair with times converted to ranges
+        """
+        if isinstance(self.time_values, bool):
+            return TimePair(self.time_type, self.time_values)
+        if self.time_type == 'week':
+            return TimePair(self.time_type, weeks_to_ranges(self.time_values))
+        return TimePair(self.time_type, days_to_ranges(self.time_values))
+
 
 def _verify_range(start: int, end: int) -> Union[int, Tuple[int, int]]:
     if start == end:
@@ -205,8 +215,28 @@ def _parse_time_pair(time_type: str, time_values: Union[bool, Sequence[str]]) ->
     raise ValidationFailedException(f'time param: {time_type} is not one of "day" or "week"')
 
 
-def parse_time_arg(key: str = "time") -> List[TimePair]:
-    return [_parse_time_pair(time_type, time_values) for [time_type, time_values] in _parse_common_multi_arg(key)]
+def parse_time_arg(key: str = "time") -> Optional[TimePair]:
+    time_pairs = [_parse_time_pair(time_type, time_values) for [time_type, time_values] in _parse_common_multi_arg(key)]
+
+    # single value
+    if len(time_pairs) == 0:
+        return None
+    if len(time_pairs) == 1:
+        return time_pairs[0]
+    
+    # make sure 'day' and 'week' aren't mixed
+    time_types = set(time_pair.time_type for time_pair in time_pairs)
+    if len(time_types) >= 2:
+        raise ValidationFailedException(f'{key}: {time_pairs} mixes "day" and "week" time types')
+
+    # merge all time pairs into one
+    merged = []    
+    for time_pair in time_pairs:
+        if time_pair.time_values == True:
+            return time_pair
+        else:
+            merged.extend(time_pair.time_values)
+    return TimePair(time_pairs[0].time_type, merged).to_ranges()
 
 
 def parse_single_time_arg(key: str) -> TimePair:
@@ -255,25 +285,26 @@ def parse_week_range_arg(key: str) -> Tuple[int, int]:
         raise ValidationFailedException(f"{key} must match YYYYWW-YYYYWW")
     return r
 
-def parse_day_or_week_arg(key: str, default_value: Optional[int] = None) -> Tuple[int, bool]:
+def parse_day_or_week_arg(key: str, default_value: Optional[int] = None) -> TimePair:
     v = request.values.get(key)
     if not v:
         if default_value is not None:
-            return default_value, guess_time_value_is_day(default_value)
+            time_type = "day" if guess_time_value_is_day(default_value) else "week"
+            return TimePair(time_type, [default_value])
         raise ValidationFailedException(f"{key} param is required")
     # format is either YYYY-MM-DD or YYYYMMDD or YYYYMM
-    is_week = len(v) == 6
+    is_week = guess_time_value_is_week(v)
     if is_week:
-        return parse_week_arg(key), False
-    return parse_day_arg(key), True
+        return TimePair("week", [parse_week_arg(key)])
+    return TimePair("day", [parse_day_arg(key)])
 
-def parse_day_or_week_range_arg(key: str) -> Tuple[Tuple[int, int], bool]:
+def parse_day_or_week_range_arg(key: str) -> TimePair:
     v = request.values.get(key)
     if not v:
         raise ValidationFailedException(f"{key} param is required")
     # format is either YYYY-MM-DD--YYYY-MM-DD or YYYYMMDD-YYYYMMDD or YYYYMM-YYYYMM
     # so if the first before the - has length 6, it must be a week
-    is_week = len(v.split('-', 2)[0]) == 6
+    is_week = guess_time_value_is_week(v.split('-', 2)[0])
     if is_week:
-        return parse_week_range_arg(key), False
-    return parse_day_range_arg(key), True
+        return TimePair("week", [parse_week_range_arg(key)])
+    return TimePair("day", [parse_day_range_arg(key)])

src/server/_query.py

@@ -20,9 +20,9 @@
 from ._db import metadata
 from ._printer import create_printer, APrinter
 from ._exceptions import DatabaseErrorException
-from ._validate import DateRange, extract_strings
+from ._validate import extract_strings
 from ._params import GeoPair, SourceSignalPair, TimePair
-from .utils import time_values_to_ranges, days_to_ranges, weeks_to_ranges
+from .utils import time_values_to_ranges, TimeValues
 
 
 def date_string(value: int) -> str:
@@ -36,7 +36,7 @@ def date_string(value: int) -> str:
 
 def to_condition(
     field: str,
-    value: Union[Tuple[str, str], str, Tuple[int, int], int],
+    value: Union[str, Tuple[int, int], int],
     param_key: str,
     params: Dict[str, Any],
     formatter=lambda x: x,
@@ -52,7 +52,7 @@ def to_condition(
 
 def filter_values(
     field: str,
-    values: Optional[Sequence[Union[Tuple[str, str], str, Tuple[int, int], int]]],
+    values: Optional[Sequence[Union[str, Tuple[int, int], int]]],
     param_key: str,
     params: Dict[str, Any],
     formatter=lambda x: x,
@@ -68,7 +68,7 @@ def filter_values(
 
 def filter_strings(
     field: str,
-    values: Optional[Sequence[Union[Tuple[str, str], str]]],
+    values: Optional[Sequence[str]],
     param_key: str,
     params: Dict[str, Any],
 ):
@@ -86,7 +86,7 @@ def filter_integers(
 
 def filter_dates(
     field: str,
-    values: Optional[Sequence[Union[Tuple[int, int], int]]],
+    values: Optional[TimeValues],
     param_key: str,
     params: Dict[str, Any],
 ):
@@ -171,32 +171,29 @@ def filter_pair(pair: SourceSignalPair, i) -> str:
     return f"({' OR '.join(parts)})"
 
 
-def filter_time_pairs(
+def filter_time_pair(
     type_field: str,
     time_field: str,
-    values: Sequence[TimePair],
+    pair: Optional[TimePair],
     param_key: str,
     params: Dict[str, Any],
 ) -> str:
     """
-    returns the SQL sub query to filter by the given time pairs
+    returns the SQL sub query to filter by the given time pair
     """
-
-    def filter_pair(pair: TimePair, i) -> str:
-        type_param = f"{param_key}_{i}t"
-        params[type_param] = pair.time_type
-        if isinstance(pair.time_values, bool) and pair.time_values:
-            return f"{type_field} = :{type_param}"
-        ranges = weeks_to_ranges(pair.time_values) if pair.is_week else days_to_ranges(pair.time_values)
-        return f"({type_field} = :{type_param} AND {filter_integers(time_field, cast(Sequence[Union[int, Tuple[int,int]]], ranges), type_param, params)})"
-
-    parts = [filter_pair(p, i) for i, p in enumerate(values)]
-
-    if not parts:
-        # something has to be selected
+    # safety path; should normally not be reached as time pairs are enforced by the API
+    if not pair:
         return "FALSE"
 
-    return f"({' OR '.join(parts)})"
+    type_param = f"{param_key}_0t"
+    params[type_param] = pair.time_type
+    if isinstance(pair.time_values, bool) and pair.time_values:
+        parts =  f"{type_field} = :{type_param}"
+    else:
+        ranges = pair.to_ranges().time_values
+        parts = f"({type_field} = :{type_param} AND {filter_integers(time_field, ranges, type_param, params)})"
+
+    return f"({parts})"
 
 
 def parse_row(
@@ -393,7 +390,7 @@ def where(self, **kvargs: Dict[str, Any]) -> "QueryBuilder":
     def where_strings(
         self,
         field: str,
-        values: Optional[Sequence[Union[Tuple[str, str], str]]],
+        values: Optional[Sequence[str]],
         param_key: Optional[str] = None,
     ) -> "QueryBuilder":
         fq_field = f"{self.alias}.{field}" if "." not in field else field
@@ -413,16 +410,6 @@ def where_integers(
         self.conditions.append(filter_integers(fq_field, values, param_key or field, self.params))
         return self
 
-    def where_dates(
-        self,
-        field: str,
-        values: Optional[Sequence[Union[Tuple[int, int], int]]],
-        param_key: Optional[str] = None,
-    ) -> "QueryBuilder":
-        fq_field = self._fq_field(field)
-        self.conditions.append(filter_dates(fq_field, values, param_key or field, self.params))
-        return self
-
     def where_geo_pairs(
         self,
         type_field: str,
@@ -463,17 +450,17 @@ def where_source_signal_pairs(
         )
         return self
 
-    def where_time_pairs(
+    def where_time_pair(
         self,
         type_field: str,
         value_field: str,
-        values: Sequence[TimePair],
+        values: Optional[TimePair],
         param_key: Optional[str] = None,
     ) -> "QueryBuilder":
         fq_type_field = self._fq_field(type_field)
         fq_value_field = self._fq_field(value_field)
         self.conditions.append(
-            filter_time_pairs(
+            filter_time_pair(
                 fq_type_field,
                 fq_value_field,
                 values,

src/server/_validate.py

@@ -3,6 +3,7 @@
 from flask import request
 
 from ._exceptions import UnAuthenticatedException, ValidationFailedException
+from .utils import TimeValues
 
 
 def resolve_auth_token() -> Optional[str]:
@@ -141,14 +142,11 @@ def extract_date(key: Union[str, Sequence[str]]) -> Optional[int]:
     return parse_date(s)
 
 
-DateRange = Union[Tuple[int, int], int]
-
-
-def extract_dates(key: Union[str, Sequence[str]]) -> Optional[List[DateRange]]:
+def extract_dates(key: Union[str, Sequence[str]]) -> Optional[TimeValues]:
     parts = extract_strings(key)
     if not parts:
         return None
-    values: List[Union[Tuple[int, int], int]] = []
+    values: TimeValues = []
 
     def push_range(first: str, last: str):
         first_d = parse_date(first)