Limit correlations:

* add test magic so as_pandas can be called

Author: dshemetov

src/server/_pandas.py

@@ -2,18 +2,19 @@
 import pandas as pd
 
 from sqlalchemy import text
+from sqlalchemy.engine.base import Engine
 
 from ._common import engine
 from ._config import MAX_RESULTS
-from ._printer import create_printer, APrinter
+from ._printer import create_printer
 from ._query import filter_fields, limit_query
 from ._exceptions import DatabaseErrorException
 
 
-def as_pandas(query: str, params: Dict[str, Any], parse_dates: Optional[Dict[str, str]] = None, limit_rows = MAX_RESULTS+1) -> pd.DataFrame:
+def as_pandas(query: str, params: Dict[str, Any], db_engine: Engine = engine, parse_dates: Optional[Dict[str, str]] = None, limit_rows = MAX_RESULTS+1) -> pd.DataFrame:
     try:
         query = limit_query(query, limit_rows)
-        return pd.read_sql_query(text(str(query)), engine, params=params, parse_dates=parse_dates)
+        return pd.read_sql_query(text(str(query)), db_engine, params=params, parse_dates=parse_dates)
     except Exception as e:
         raise DatabaseErrorException(str(e))
 

src/server/_query.py

@@ -233,13 +233,13 @@ def parse_result(
 
 
 def limit_query(query: str, limit: int) -> str:
-    # limit rows + 1 for detecting whether we would have more
     full_query = f"{query} LIMIT {limit}"
     return full_query
 
 
 def run_query(p: APrinter, query_tuple: Tuple[str, Dict[str, Any]]):
     query, params = query_tuple
+    # limit rows + 1 for detecting whether we would have more
     full_query = text(limit_query(query, p.remaining_rows + 1))
     app.logger.info("full_query: %s, params: %s", full_query, params)
     return db.execution_options(stream_results=True).execute(full_query, **params)

tests/server/test_pandas.py

@@ -6,14 +6,13 @@
 import unittest
 
 import pandas as pd
-from pandas.core.frame import DataFrame
-from sqlalchemy import text
-import mysql.connector
+from sqlalchemy import create_engine
 
 # from flask.testing import FlaskClient
 from delphi_utils import Nans
+from delphi.epidata.server.main import app
 from delphi.epidata.server._pandas import as_pandas
-from delphi.epidata.server._query import limit_query
+from delphi.epidata.server._query import QueryBuilder
 
 # py3tester coverage target
 __test_target__ = "delphi.epidata.server._query"
@@ -111,27 +110,26 @@ class UnitTests(unittest.TestCase):
 
     def setUp(self):
         """Perform per-test setup."""
+        app.config["TESTING"] = True
+        app.config["WTF_CSRF_ENABLED"] = False
+        app.config["DEBUG"] = False
 
         # connect to the `epidata` database and clear the `covidcast` table
-        cnx = mysql.connector.connect(user="user", password="pass", host="delphi_database_epidata", database="epidata")
-        cur = cnx.cursor()
-        cur.execute("truncate table covidcast")
-        cur.execute('update covidcast_meta_cache set timestamp = 0, epidata = ""')
-        cnx.commit()
-        cur.close()
+        engine = create_engine('mysql://user:pass@delphi_database_epidata/epidata')
+        cnx = engine.connect()
+        cnx.execute("truncate table covidcast")
+        cnx.execute('update covidcast_meta_cache set timestamp = 0, epidata = ""')
 
         # make connection and cursor available to test cases
         self.cnx = cnx
-        self.cur = cnx.cursor()
 
     def tearDown(self):
         """Perform per-test teardown."""
-        self.cur.close()
         self.cnx.close()
 
     def _insert_rows(self, rows: Iterable[CovidcastRow]):
         sql = ",\n".join((str(r) for r in rows))
-        self.cur.execute(
+        self.cnx.execute(
             f"""
             INSERT INTO
                 `covidcast` (`id`, `source`, `signal`, `time_type`, `geo_type`,
@@ -143,7 +141,6 @@ def _insert_rows(self, rows: Iterable[CovidcastRow]):
             {sql}
             """
         )
-        self.cnx.commit()
         return rows
 
     def _rows_to_df(self, rows: Iterable[CovidcastRow]) -> pd.DataFrame:
@@ -160,17 +157,12 @@ def test_as_pandas(self):
         rows = [CovidcastRow(time_value=20200401 + i, value=float(i)) for i in range(10)]
         self._insert_rows(rows)
 
-        with self.subTest("simple"):
-            query = """select * from `covidcast`"""
-            params = {}
-            parse_dates = None
-            engine = self.cnx
-            df = pd.read_sql_query(str(query), engine, params=params, parse_dates=parse_dates)
-            df = df.astype({"is_latest_issue": bool, "is_wip": bool})
+        with app.test_request_context('/correlation'):
+            q = QueryBuilder("covidcast", "t")
+
+            df = as_pandas(str(q), params={}, db_engine=self.cnx, parse_dates=None).astype({"is_latest_issue": bool, "is_wip": bool})
             expected_df = self._rows_to_df(rows)
             pd.testing.assert_frame_equal(df, expected_df)
-            query = limit_query(query, 5)
-            df = pd.read_sql_query(str(query), engine, params=params, parse_dates=parse_dates)
-            df = df.astype({"is_latest_issue": bool, "is_wip": bool})
+            df = as_pandas(str(q), params={}, db_engine=self.cnx, parse_dates=None, limit_rows=5).astype({"is_latest_issue": bool, "is_wip": bool})
             expected_df = self._rows_to_df(rows[:5])
             pd.testing.assert_frame_equal(df, expected_df)