From da558b1cf5012eb48e48000a6619e249cf05192e Mon Sep 17 00:00:00 2001 From: Ted Kaemming Date: Wed, 14 Aug 2019 12:13:05 -0700 Subject: [PATCH 1/3] fix(reader): Correctly handle totals row --- snuba/reader.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/snuba/reader.py b/snuba/reader.py index 55b7840e3d8..e8d9a298b98 100644 --- a/snuba/reader.py +++ b/snuba/reader.py @@ -4,7 +4,7 @@ import re from abc import ABC, abstractmethod from datetime import datetime -from typing import TYPE_CHECKING, Any, Mapping, Optional, Sequence +from typing import TYPE_CHECKING, Any, Mapping, MutableMapping, Optional, Sequence from dateutil.tz import tz @@ -13,12 +13,13 @@ from mypy_extensions import TypedDict Column = TypedDict("Column", {"name": str, "type": str}) + Row = MutableMapping[str, Any] Result = TypedDict( "Result", { "meta": Sequence[Column], - "data": Sequence[Mapping[str, Any]], - "totals": Mapping[str, Any], + "data": Sequence[Row], + "totals": Row, }, total=False, ) @@ -48,10 +49,10 @@ def transform_date_columns(result: Result) -> Result: """ for col in result["meta"]: if DATETIME_TYPE_RE.match(col["type"]): - for row in itertools.chain(result["data"], result.get("totals", [])): + for row in itertools.chain(result["data"], [result.get("totals", {})]): row[col["name"]] = row[col["name"]].replace(tzinfo=tz.tzutc()) elif DATE_TYPE_RE.match(col["type"]): - for row in itertools.chain(result["data"], result.get("totals", [])): + for row in itertools.chain(result["data"], [result.get("totals", {})]): row[col["name"]] = datetime( *(row[col["name"]].timetuple()[:6]) ).replace(tzinfo=tz.tzutc()) From 389cbd087353377767fffad175c69720393b52dc Mon Sep 17 00:00:00 2001 From: Ted Kaemming Date: Wed, 14 Aug 2019 12:25:36 -0700 Subject: [PATCH 2/3] be better at programming --- snuba/reader.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/snuba/reader.py b/snuba/reader.py index e8d9a298b98..a0184da1d3a 100644 --- a/snuba/reader.py +++ b/snuba/reader.py @@ -47,15 +47,23 @@ def transform_date_columns(result: Result) -> Result: Convert timezone-naive date and datetime values into timezone aware datetimes. """ + + def iterate_rows(): + if 'totals' in result: + return itertools.chain(result['data'], [result['totals']]) + else: + return iter(result["data"]) + for col in result["meta"]: if DATETIME_TYPE_RE.match(col["type"]): - for row in itertools.chain(result["data"], [result.get("totals", {})]): + for row in iterate_rows(): row[col["name"]] = row[col["name"]].replace(tzinfo=tz.tzutc()) elif DATE_TYPE_RE.match(col["type"]): - for row in itertools.chain(result["data"], [result.get("totals", {})]): + for row in iterate_rows(): row[col["name"]] = datetime( *(row[col["name"]].timetuple()[:6]) ).replace(tzinfo=tz.tzutc()) + return result From a1ac6e509927a84681ee9625adc97800fc1c0f4d Mon Sep 17 00:00:00 2001 From: Ted Kaemming Date: Wed, 14 Aug 2019 12:26:44 -0700 Subject: [PATCH 3/3] format --- snuba/reader.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/snuba/reader.py b/snuba/reader.py index a0184da1d3a..35c9e52e04f 100644 --- a/snuba/reader.py +++ b/snuba/reader.py @@ -16,11 +16,7 @@ Row = MutableMapping[str, Any] Result = TypedDict( "Result", - { - "meta": Sequence[Column], - "data": Sequence[Row], - "totals": Row, - }, + {"meta": Sequence[Column], "data": Sequence[Row], "totals": Row}, total=False, ) @@ -49,8 +45,8 @@ def transform_date_columns(result: Result) -> Result: """ def iterate_rows(): - if 'totals' in result: - return itertools.chain(result['data'], [result['totals']]) + if "totals" in result: + return itertools.chain(result["data"], [result["totals"]]) else: return iter(result["data"])