feat: simple running version

Author: sgratzl

mypy.ini

@@ -0,0 +1,2 @@
+[mypy]
+plugins = sqlmypy

requirements.dev.txt

@@ -1,2 +1,4 @@
 invoke>=1.4.1
-black>=20.8b1
+black>=20.8b1
+sqlalchemy-stubs>=0.3
+mypy>=0.790

requirements.txt

@@ -1,4 +1,4 @@
 Flask==1.1.2
 SQLAlchemy==1.3.22
-PyMySQL==0.10.1
+mysqlclient==2.0.2
 python-dotenv==0.15.0

src/server/_analytics.py

@@ -1,5 +1,5 @@
 from sqlalchemy import Table, Column, String, Integer, DateTime, func
-from ._db import metadata
+from ._db import metadata, TABLE_OPTIONS
 from flask import request
 from ._common import db
 
@@ -13,6 +13,7 @@
     Column("source", String(32)),
     Column("result", Integer),
     Column("num_rows", Integer),
+    **TABLE_OPTIONS
 )
 
 

src/server/_db.py

@@ -4,3 +4,10 @@
 
 engine: Engine = create_engine(SQLALCHEMY_DATABASE_URI, **SQLALCHEMY_ENGINE_OPTIONS)
 metadata = MetaData(bind=engine)
+
+TABLE_OPTIONS = dict(
+    mysql_engine="InnoDB",
+    # mariadb_engine="InnoDB",
+    mysql_charset="utf8mb4",
+    # mariadb_charset="utf8",
+)

src/server/_printer.py

@@ -1,8 +1,17 @@
-from flask import request, jsonify, abort, make_response, request
+from flask import (
+    request,
+    jsonify,
+    abort,
+    make_response,
+    request,
+    Response,
+    stream_with_context,
+)
 from werkzeug.exceptions import HTTPException
 from flask.json import dumps
 from ._analytics import record_analytics
-from typing import Dict
+from typing import Dict, Iterable, Any, Union, Optional, List
+
 
 MAX_RESULTS = 1000
 
@@ -48,13 +57,12 @@ def __init__(self):
 class APrinter:
     count: int = 0
     result: int = -1
-    _began = False
-    _use_status_codes = True
-    endpoint: str
 
-    def __init__(self, endpoint: str, use_status_codes=True):
-        self.endpoint = endpoint
-        self._use_status_codes = use_status_codes
+    def make_response(self, gen):
+        return Response(
+            gen,
+            mimetype="application/json",
+        )
 
     def print_non_standard(self, data):
         """
@@ -67,163 +75,102 @@ def print_non_standard(self, data):
 
     @property
     def remaining_rows(self) -> int:
-        return MAX_RESULTS - self._count
-
-    def __enter__(self):
-        self._begin()
-        return self
+        return MAX_RESULTS - self.count
 
-    def _begin(self):
-        if self._began:
-            return
-        self._began = True
+    def begin(self):
         self.result = -2  # no result
-        self._begin_impl()
+        return self._begin_impl()
 
     def _begin_impl(self):
         # hook
-        pass
+        return None
 
-    def __call__(self, row: Dict):
-        if not self._began:
-            self._begin()
+    def __call__(self, row: Dict) -> Optional[Union[str, bytes]]:
         first = self.count == 0
         if self.count >= MAX_RESULTS:
             # hit the limit
             self.result = 2
-            return
+            return None
         if first:
             self.result = 1  # at least one row
-        self._print_row(first, row)
         self.count += 1
+        return self._print_row(first, row)
 
-    def _print_row_impl(self, first: bool, row: Dict):
-        pass
+    def _print_row(self, first: bool, row: Dict) -> Optional[Union[str, bytes]]:
+        return None
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        if not self._began:
-            return
-        # record_analytics($this->endpoint, $this->result, $this->count)
-        self._end_impl()
+    def end(self) -> Optional[Union[str, bytes]]:
+        record_analytics(self.result, self.count)
+        return self._end_impl()
 
     def _end_impl(self):
         # hook
-        pass
+        return None
+
+
+class ClassicPrinter(APrinter):
+    """
+    a printer class writing in the classic epidata format
+    """
 
+    def _begin_impl(self):
+        return '{ "epidata": ['
 
-# interface IRowPrinter {
-#   /**
-#    * returns the number of possible remaining rows to fetch, if < 0 it means no limit
-#    */
-#   public function remainingRows(): int;
-#   /**
-#    * starts printing rows, can be called multiple times without harm
-#    */
-#   public function begin();
-#   /**
-#    * print a specific row
-#    */
-#   public function printRow(array &$row);
-#   /**
-#    * finish writing rows needs to be called once
-#    */
-#   public function end();
-# }
-
-# class CollectRowPrinter implements IRowPrinter {
-#   public array $data = [];
-
-#   public function remainingRows(): int {
-#     global $MAX_RESULTS;
-#     return $MAX_RESULTS - count($this->data);
-#   }
-
-#   public function begin() {
-#     // dummy
-#   }
-#   public function printRow(array &$row) {
-#     array_push($this->data, $row);
-#   }
-#   public function end() {
-#     // dummy
-#   }
-# }
-
-# /**
-#  * a printer class writing in the classic epidata format
-#  */
-# class ClassicPrinter extends APrinter {
-
-#   function __construct(string $endpoint) {
-#     parent::__construct($endpoint, FALSE);
-#   }
-
-#   protected function beginImpl() {
-#     header('Content-Type: application/json');
-#     echo '{ "epidata": [';
-#   }
-
-#   protected function printRowImpl(bool $first, array &$row) {
-#     if (!$first) {
-#       echo ',';
-#     }
-#     echo json_encode($row);
-#   }
-
-#   protected function endImpl() {
-#     $message = $this->count == 0 ? 'no results' : ($this->result == 2 ? 'too many results, data truncated' : 'success');
-#     $messageEncoded = json_encode($message);
-#     echo "], \"result\": {$this->result}, \"message\": {$messageEncoded} }";
-#   }
-# }
-
-# /**
-#  * a printer class writing a tree by the given grouping criteria as the first element in the epidata array
-#  */
-# class ClassicTreePrinter extends ClassicPrinter {
-#   private array $tree = [];
-#   private string $group;
-
-#   function __construct(string $endpoint, string $group) {
-#     parent::__construct($endpoint);
-#     $this->group = $group;
-#   }
-
-#   protected function printRowImpl(bool $first, array &$row) {
-#     $group = isset($row[$this->group]) ? $row[$this->group] : '';
-#     unset($row[$this->group]);
-#     if (isset($this->tree[$group])) {
-#       array_push($this->tree[$group], $row);
-#     } else {
-#       $this->tree[$group] = [$row];
-#     }
-#   }
-
-#   private function printTree() {
-#     if (count($this->tree) == 0) {
-#       echo '{}'; // force object style
-#     } else {
-#       echo json_encode($this->tree);
-#     }
-#     // clean up
-#     $this->tree = [];
-#   }
-
-#   protected function endImpl() {
-#     $this->printTree();
-#     parent::endImpl();
-#   }
-# }
+    def _print_row(self, first: bool, row: Dict):
+        sep = "," if not first else ""
+        return f"{sep}{dumps(row)}"
+
+    def _end_impl(self):
+        message = "success"
+        if self.count == 0:
+            message = "no results"
+        elif self.result == 2:
+            message = "too many results, data truncated"
+        return f'], "result": {self.result}, "message": {dumps(message)} }}'
+
+
+class ClassicTreePrinter(ClassicPrinter):
+    """
+    a printer class writing a tree by the given grouping criteria as the first element in the epidata array
+    """
+
+    group: str
+    _tree: Dict[str, List[Dict]] = dict()
+
+    def __init__(self, group: str):
+        super(ClassicTreePrinter, self).__init__()
+        self.group = group
+        self._tree = dict()
+
+    def _print_row(self, first: bool, row: Dict):
+        group = row.get(self.group, "")
+        del row[self.group]
+        if group in self._tree:
+            self._tree[group].append(row)
+        else:
+            self._tree[group] = [row]
+        return None
+
+    def _end_impl(self):
+        tree = dumps(self._tree)
+        self._tree = dict()
+        r = super(ClassicTreePrinter, self)._end_impl()
+        return f"{tree}{r}"
 
 
 class CSVPrinter(APrinter):
     """
     a printer class writing in a CSV file
     """
 
+    def make_response(self, gen):
+        return Response(
+            gen,
+            mimetype="text/csv; charset=utf8",
+            headers={"Content-Disposition": "attachment; filename=epidata.csv"},
+        )
+
     def _begin_impl(self):
-        # header('Content-Type: text/csv; charset=utf8');
-        # header('Content-Disposition: attachment; filename=epidata.csv');
         pass
 
     def _print_row_impl(self, first: bool, row: Dict):
@@ -246,43 +193,58 @@ class JSONPrinter(APrinter):
     """
 
     def _begin_impl(self):
-        # TODO header('Content-Type: application/json')
-        # echo '['
-        pass
+        return "["
 
-    def _print_row_impl(self, first: bool, row: Dict):
+    def _print_row(self, first: bool, row: Dict):
         sep = "," if not first else ""
         return f"{sep}{dumps(row)}"
 
     def _end_impl(self):
-        # echo ']'
-        pass
+        return "]"
 
 
 class JSONLPrinter(APrinter):
     """
     a printer class writing in JSONLines format
     """
 
-    def _begin_impl(self):
-        # TODO
-        # there is no official mime type for json lines
-        # header('Content-Type: text/plain; charset=utf8');
-        pass
+    def make_response(self, gen):
+        return Response(gen, mimetype=" text/plain; charset=utf8")
 
-    def _print_row_impl(self, first: bool, row: Dict):
+    def _print_row(self, first: bool, row: Dict):
         # each line is a JSON file with a new line to separate them
         return f"{dumps(row)}\n"
 
 
-def create_printer(endpoint: str):
+def create_printer():
     format = request.values.get("format", "classic")
-    # if format == "tree":
-    #     return ClassicTreePrinter(endpoint, "signal")
-    # if format == "json":
-    return JSONPrinter(endpoint)
-    # if format == "csv":
-    #     return CSVPrinter(endpoint)
-    # if format == "jsonl":
-    #     return JSONLPrinter(endpoint)
-    # return ClassicPrinter(endpoint)
+    if format == "tree":
+        return ClassicTreePrinter("signal")
+    if format == "json":
+        return JSONPrinter()
+    if format == "csv":
+        return CSVPrinter()
+    if format == "jsonl":
+        return JSONLPrinter()
+    return ClassicPrinter()
+
+
+def send_rows(generator: Iterable[Dict[str, Any]]):
+    printer = create_printer()
+
+    def gen():
+        r = printer.begin()
+        if r is not None:
+            yield r
+        for row in generator:
+            r = printer(row)
+            if r is not None:
+                yield r
+        r = printer.end()
+        if r is not None:
+            yield r
+
+    return printer.make_response(stream_with_context(gen()))
+
+
+# , execution_options={"stream_results": True}