diff --git a/corrlib/find.py b/corrlib/find.py index e4ee735..4c51e05 100644 --- a/corrlib/find.py +++ b/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: - 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: +def _db_lookup(db: Path, ensemble: str, correlator_name: str, code: str, project: Optional[str]=None, parameters: Optional[str]=None, + created_before: Optional[str]=None, created_after: Optional[Any]=None, updated_before: Optional[Any]=None, updated_after: Optional[Any]=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) - if Any([created_before, created_after, updated_before, updated_after]): - results = _time_filter(results, created_before, created_after, updated_before, updated_after) + 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 code == "sfcf": results = sfcf_filter(results, **kwargs) elif code == "openQCD": diff --git a/tests/find_test.py b/tests/find_test.py index da1bfc1..e0730e9 100644 --- a/tests/find_test.py +++ b/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) - -