add test for db lookup

corrlib/find.py

@@ -8,6 +8,7 @@ from .tools import k2m, get_db_file
 from .tracker import get
 from typing import Any, Optional
 from pathlib import Path
+import datetime as dt
 
 
 def _project_lookup_by_alias(db: Path, alias: str) -> str:
@@ -62,8 +63,37 @@ def _project_lookup_by_id(db: Path, uuid: str) -> list[tuple[str, str]]:
     return results
 
 
-def _db_lookup(db: Path, ensemble: str, correlator_name: str, code: str, project: Optional[str]=None,  parameters: Optional[str]=None,
+def _time_filter(results: pd.DataFrame,  created_before: Optional[str]=None, created_after: Optional[Any]=None, updated_before: Optional[Any]=None, updated_after: Optional[Any]=None) -> pd.DataFrame:
-               created_before: Optional[str]=None, created_after: Optional[Any]=None, updated_before: Optional[Any]=None, updated_after: Optional[Any]=None) -> pd.DataFrame:
+    drops = []
+    for ind in len(results):
+        result = results.iloc[ind]
+        created_at = dt.datetime.fromisoformat(result['created_at'])
+        updated_at = dt.datetime.fromisoformat(result['updated_at'])
+
+        if created_before is not None:
+            created_before = dt.datetime.fromisoformat(created_before)
+            if created_before < created_at:
+                drops.append(ind)
+                continue
+        if created_after is not None:
+            created_after = dt.datetime.fromisoformat(created_after)
+            if created_before > created_at:
+                drops.append(ind)
+                continue
+        if updated_before is not None:
+            updated_before = dt.datetime.fromisoformat(updated_before)
+            if updated_before < updated_at:
+                drops.append(ind)
+                continue
+        if updated_after is not None:
+            updated_after = dt.datetime.fromisoformat(updated_after)
+            if updated_after > updated_at:
+                drops.append(ind)
+                continue
+        return results.drop(drops)
+
+
+def _db_lookup(db: Path, ensemble: str, correlator_name: str, code: str, project: Optional[str]=None,  parameters: Optional[str]=None) -> pd.DataFrame:
     """
     Look up a correlator record in the database by the data given to the method.
 
@@ -105,14 +135,6 @@ def _db_lookup(db: Path, ensemble: str, correlator_name: str, code: str, project
         search_expr += f" AND code = '{code}'"
     if parameters:
         search_expr += f" AND parameters = '{parameters}'"
-    if created_before:
-        search_expr += f" AND created_at < '{created_before}'"
-    if created_after:
-        search_expr += f" AND created_at > '{created_after}'"
-    if updated_before:
-        search_expr += f" AND updated_at < '{updated_before}'"
-    if updated_after:
-        search_expr += f" AND updated_at > '{updated_after}'"
     conn = sqlite3.connect(db)
     results = pd.read_sql(search_expr, conn)
     conn.close()
@@ -236,7 +258,9 @@ def find_record(path: Path, ensemble: str, correlator_name: str, code: str, proj
     if code not in codes:
         raise ValueError("Code " + code + "unknown, take one of the following:" + ", ".join(codes))
     get(path, db_file)
-    results = _db_lookup(db, ensemble, correlator_name,code, project, parameters=parameters, created_before=created_before, created_after=created_after, updated_before=updated_before, updated_after=updated_after)
+    results = _db_lookup(db, ensemble, correlator_name,code, project, parameters=parameters)
+    if Any([created_before, created_after, updated_before, updated_after]):
+        results = _time_filter(results, created_before, created_after, updated_before, updated_after)
     if code == "sfcf":
         results = sfcf_filter(results, **kwargs)
     elif code == "openQCD":

tests/find_test.py

@@ -54,3 +54,42 @@ def test_find_lookup_by_id(tmp_path: Path) -> None:
     assert owner == result[3]
     assert code == result[4]
 
+
+def test_db_lookup(tmp_path: Path) -> None:
+    db = make_sql(tmp_path)
+    conn = sqlite3.connect(db)
+    c = conn.cursor()
+
+    corr = "f_A"
+    ensemble = "SF_A"
+    code = "openQCD"
+    meas_path = "archive/SF_A/f_A/Project_A.json.gz::asdfasdfasdf"
+    uuid = "Project_A"
+    pars = "{par_A: 3.0, par_B: 5.0}"
+    parameter_file = "projects/Project_A/myinput.in"
+    c.execute("INSERT INTO backlogs (name, ensemble, code, path, project, parameters, parameter_file, created_at, updated_at) VALUES (?, ?, ?, ?, ?, ?, ?, datetime('now'), datetime('now'))",
+                                                (corr, ensemble, code, meas_path, uuid, pars, parameter_file))
+    conn.commit()
+    conn.close()
+
+    results = find._db_lookup(db, ensemble, corr, code)
+    assert len(results) == 1
+    results = find._db_lookup(db, "SF_B", corr, code)
+    assert results.empty
+    results = find._db_lookup(db, ensemble, "g_A", code)
+    assert results.empty
+    results = find._db_lookup(db, ensemble, corr, "sfcf")
+    assert results.empty
+    results = find._db_lookup(db, ensemble, corr, code, project = "Project_A")
+    assert len(results) == 1
+    results = find._db_lookup(db, ensemble, corr, code, project = "Project_B")
+    assert results.empty
+    results = find._db_lookup(db, ensemble, corr, code, parameters = pars)
+    assert len(results) == 1
+    results = find._db_lookup(db, ensemble, corr, code, parameters = "{par_A: 3.0, par_B: 4.0}")
+    assert results.empty
+    #results = find._db_lookup(db, ensemble, corr, code, project,  parameters, created_before, created_after, updated_before, updated_after)
+
+    #results = find._db_lookup(db, ensemble, corr, code, project,  parameters, created_before, created_after, updated_before, updated_after)
+
+