feat(reader): Add ClickHouse HTTPReader

snuba/clickhouse/http.py

@@ -1,12 +1,18 @@
+import json
 import re
+from typing import Callable, Iterable, Mapping, MutableMapping, Optional
 from urllib.parse import urlencode
-from typing import Callable, Iterable
 
+from dateutil.parser import parse as dateutil_parse
 from urllib3.connectionpool import HTTPConnectionPool
 from urllib3.exceptions import HTTPError
+from urllib3.response import HTTPResponse
 
 from snuba.clickhouse.errors import ClickhouseError
+from snuba.clickhouse.native import transform_date, transform_datetime
+from snuba.clickhouse.query import ClickhouseQuery
 from snuba.datasets.schemas.tables import TableSchema
+from snuba.reader import Reader, Result, build_result_transformer
 from snuba.writer import BatchWriter, WriterTableRow
 
 
@@ -16,6 +22,21 @@
 )
 
 
+def raise_for_error_response(response: HTTPResponse) -> None:
+    if response.status != 200:
+        # XXX: This should be switched to just parse the JSON body after
+        # https://github.com/yandex/ClickHouse/issues/6272 is available.
+        content = response.data.decode("utf8")
+        details = CLICKHOUSE_ERROR_RE.match(content)
+        if details is not None:
+            code, type, message = details.groups()
+            raise ClickhouseError(int(code), message)
+        else:
+            raise HTTPError(
+                f"Received unexpected {response.status} response: {content}"
+            )
+
+
 class HTTPBatchWriter(BatchWriter):
     def __init__(
         self,
@@ -73,15 +94,80 @@ def write(self, rows: Iterable[WriterTableRow]):
             chunked=True,
         )
 
-        if response.status != 200:
-            # XXX: This should be switched to just parse the JSON body after
-            # https://github.com/yandex/ClickHouse/issues/6272 is available.
-            content = response.data.decode("utf8")
-            details = CLICKHOUSE_ERROR_RE.match(content)
-            if details is not None:
-                code, type, message = details.groups()
-                raise ClickhouseError(int(code), message)
-            else:
-                raise HTTPError(
-                    f"Received unexpected {response.status} response: {content}"
-                )
+        raise_for_error_response(response)
+
+
+def parse_and_transform_date(value: str) -> str:
+    return transform_date(dateutil_parse(value))
+
+
+def parse_and_transform_datetime(value: str) -> str:
+    return transform_datetime(dateutil_parse(value))
+
+
+transform_column_types = build_result_transformer(
+    [
+        (re.compile(r"^Date(\(.+\))?$"), parse_and_transform_date),
+        (re.compile(r"^DateTime(\(.+\))?$"), parse_and_transform_datetime),
+    ]
+)
+
+
+class HTTPReader(Reader[ClickhouseQuery]):
+    def __init__(
+        self, host: str, port: int, settings: Optional[Mapping[str, str]] = None
+    ):
+        if settings is not None:
+            assert "query_id" not in settings, "query_id cannot be passed as a setting"
+
+        self.__pool = HTTPConnectionPool(host, port)
+
+        self.__default_settings: MutableMapping[str, str] = (
+            {**settings} if settings is not None else {}
+        )
+
+        if "output_format_json_quote_64bit_integers" not in self.__default_settings:
+            self.__default_settings["output_format_json_quote_64bit_integers"] = "0"
+
+    def execute(
+        self,
+        query: ClickhouseQuery,
+        settings: Optional[Mapping[str, str]] = None,
+        query_id: Optional[str] = None,
+        with_totals: bool = False,  # NOTE: unnecessary with FORMAT JSON
+    ) -> Result:
+        query_settings: MutableMapping[str, str] = (
+            {**settings} if settings is not None else {}
+        )
+
+        # XXX: mypy won't allow redefining ``settings`` as mutable, so delete
+        # the original variable to avoid accidentally referencing ``settings``
+        # instead of ``query_settings``.
+        del settings
+
+        assert (
+            "query_id" not in query_settings
+        ), "query_id cannot be passed as a setting"
+
+        if query_id is not None:
+            query_settings["query_id"] = query_id
+
+        response = self.__pool.urlopen(
+            "POST",
+            "/?" + urlencode({**self.__default_settings, **query_settings}),
+            headers={"Connection": "keep-alive", "Accept-Encoding": "gzip,deflate"},
+            body=query.format_sql("JSON"),
+        )
+
+        raise_for_error_response(response)
+
+        result = json.loads(response.data.decode("utf-8"))
+
+        # Remove any extra keys that are not part of the Result data structure.
+        for k in [*result.keys()]:
+            if k not in {"meta", "data", "totals"}:
+                del result[k]
+
+        transform_column_types(result)
+
+        return result

tests/clickhouse/test_reader.py

@@ -0,0 +1,60 @@
+from typing import Sequence, Tuple
+
+import pytest
+
+from snuba import settings
+from snuba.clickhouse.errors import ClickhouseError
+from snuba.clickhouse.http import HTTPReader
+from snuba.clickhouse.native import NativeDriverReader
+from snuba.clickhouse.query import ClickhouseQuery
+from snuba.environment import clickhouse_ro
+from snuba.reader import Reader
+
+
+class SimpleClickhouseQuery(ClickhouseQuery):
+    def __init__(self, columns: Sequence[Tuple[str, str]]) -> None:
+        self.__columns = columns
+
+    def _format_query_impl(self) -> str:
+        columns = ", ".join(f"{value} as {alias}" for alias, value in self.__columns)
+        return f"SELECT {columns}"
+
+
+@pytest.mark.parametrize(
+    "reader",
+    [
+        NativeDriverReader(clickhouse_ro),
+        HTTPReader(settings.CLICKHOUSE_HOST, settings.CLICKHOUSE_HTTP_PORT),
+    ],
+)
+def test_reader(reader: Reader[ClickhouseQuery]) -> None:
+    assert reader.execute(
+        SimpleClickhouseQuery(
+            [
+                ("datetime", "toDateTime('2020-01-02 03:04:05')"),
+                ("date", "toDate('2020-01-02')"),
+                ("int64", "toInt64(1)"),
+                ("uuid", "toUUID('00000000-0000-4000-8000-000000000000')"),
+            ]
+        )
+    ) == {
+        "meta": [
+            {"name": "datetime", "type": "DateTime"},
+            {"name": "date", "type": "Date"},
+            {"name": "int64", "type": "Int64"},
+            {"name": "uuid", "type": "UUID"},
+        ],
+        "data": [
+            {
+                "date": "2020-01-02T00:00:00+00:00",
+                "datetime": "2020-01-02T03:04:05+00:00",
+                "int64": 1,
+                "uuid": "00000000-0000-4000-8000-000000000000",
+            }
+        ],
+    }
+
+    with pytest.raises(ClickhouseError) as e:
+        reader.execute(SimpleClickhouseQuery([("invalid", '"')]))
+
+    assert e.value.code == 62