corrlib/find.py

@@ -8,7 +8,6 @@ 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:
@@ -63,37 +62,8 @@ def _project_lookup_by_id(db: Path, uuid: str) -> list[tuple[str, str]]:
     return results
 
 
-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:
+def _db_lookup(db: Path, ensemble: str, correlator_name: str, code: str, project: Optional[str]=None,  parameters: Optional[str]=None,
-    drops = []
+               created_before: Optional[str]=None, created_after: Optional[Any]=None, updated_before: Optional[Any]=None, updated_after: Optional[Any]=None) -> pd.DataFrame:
-    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.
 
@@ -135,6 +105,14 @@ 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()
@@ -258,9 +236,7 @@ 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)
+    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)
-    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,42 +54,3 @@ 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)
-
-