Improve publication fetching for datatables.

3 files changed, 60 insertions, 40 deletions

diff --git a/uploader/publications/datatables.py b/uploader/publications/datatables.py
new file mode 100644
index 0000000..e07fafd
--- /dev/null
+++ b/uploader/publications/datatables.py
@@ -0,0 +1,52 @@
+"""Fetch data for datatables."""
+import logging
+from typing import Optional
+
+from MySQLdb.cursors import DictCursor
+
+from gn_libs.mysqldb import Connection, debug_query
+
+logger = logging.getLogger(__name__)
+
+def fetch_publications(
+        conn: Connection,
+        search: Optional[str] = None,
+        offset: int = 0,
+        limit: int = -1
+) -> tuple[dict, int, int, int]:
+    """Fetch publications from the database."""
+    _query = "SELECT * FROM Publication"
+    _count_query = "SELECT COUNT(*) FROM Publication"
+    _params = None
+    _where_clause = ""
+    _limit_clause = ""
+    if search is not None and bool(search):
+        _where_clause = ("WHERE PubMed_ID LIKE %s "
+                      "OR Authors LIKE %s "
+                      "OR Title LIKE %s")
+        _params = (f"%{search}%",) * 3
+
+    if limit > 0:
+        _limit_clause = f"LIMIT {limit} OFFSET {offset}"
+
+    with conn.cursor(cursorclass=DictCursor) as cursor:
+        cursor.execute("SELECT COUNT(*) FROM Publication")
+        _total_rows = int(cursor.fetchone()["COUNT(*)"])
+
+        cursor.execute(f"{_count_query} {_where_clause}", _params)
+        debug_query(cursor, logger)
+        _result = cursor.fetchone()
+        _total_filtered = int(_result["COUNT(*)"] if bool(_result) else 0)
+
+        cursor.execute(f"{_query} {_where_clause} {_limit_clause}", _params)
+        debug_query(cursor, logger)
+        _current_filtered = tuple(
+            {**dict(row), "index": idx}
+            for idx, row
+            in enumerate(cursor.fetchall(), start=offset+1))
+
+        return (
+            _current_filtered,
+            len(_current_filtered),
+            _total_filtered,
+            _total_rows)
diff --git a/uploader/publications/models.py b/uploader/publications/models.py
index f4e8d33..b199991 100644
--- a/uploader/publications/models.py
+++ b/uploader/publications/models.py
@@ -69,32 +69,6 @@ def update_publications(conn: Connection , publications: tuple[dict, ...]) -> tu
     return tuple()
 
 
-def fetch_publications(
-        conn: Connection,
-        search: Optional[str] = None,
-        offset: int = 0,
-        limit: int = -1
-) -> Iterable[dict]:
-    """Fetch publications from the database."""
-    _query = "SELECT * FROM Publication"
-    _params = None
-    if search is not None and bool(search):
-        _query = (f"{_query} "
-                  "WHERE PubMed_ID LIKE %s "
-                  "OR Authors LIKE %s "
-                  "OR Title LIKE %s")
-        _params = (f"%{search}%",) * 3
-
-    if limit > 0:
-        _query = (f"{_query} LIMIT {limit} OFFSET {offset}")
-
-    with conn.cursor(cursorclass=DictCursor) as cursor:
-        cursor.execute(_query, _params)
-        debug_query(_query, logger)
-        for row in cursor.fetchall():
-            yield dict(row)
-
-
 def fetch_publication_by_id(conn: Connection, publication_id: int) -> dict:
     """Fetch a specific publication from the database."""
     with conn.cursor(cursorclass=DictCursor) as cursor:
diff --git a/uploader/publications/views.py b/uploader/publications/views.py
index 8a65ff2..137052b 100644
--- a/uploader/publications/views.py
+++ b/uploader/publications/views.py
@@ -15,11 +15,12 @@ from flask import (
 from uploader.authorisation import require_login
 
 from .models import (
-    fetch_publications,
     fetch_publication_by_id,
     create_new_publications,
     fetch_publication_phenotypes)
 
+from .datatables import fetch_publications
+
 from gn_libs.debug import __pk__
 
 pubbp = Blueprint("publications", __name__)
@@ -44,23 +45,16 @@ def list_publications():
     _search = request.args["search[value]"]
     with (database_connection(app.config["SQL_URI"]) as conn,
           conn.cursor(cursorclass=DictCursor) as cursor):
-        cursor.execute("SELECT COUNT(*) FROM Publication")
-        _totalrows = int(cursor.fetchone()["COUNT(*)"])
-        _publications = tuple({
-            **row, "index": idx
-        } for idx,row in enumerate(
-            fetch_publications(
-                conn,
-                _search,
-                offset=_length * (_page-1),
-                limit=_length),
-            start=(_start + 1)))
+        _publications, _current_rows, _totalfiltered, _totalrows = fetch_publications(
+            conn,
+            _search,
+            offset=_start,
+            limit=_length)
 
         return json.dumps({
             "draw": _page,
             "recordsTotal": _totalrows,
-            "recordsFiltered": (
-                len(_publications) if bool(_search) else _totalrows),
+            "recordsFiltered": _totalfiltered,
             "publications": _publications,
             "status": "success"
         })