implement find functionality

backlogger/find.py

@@ -0,0 +1,91 @@
+import sqlite3
+import datalad.api as dl
+import os
+import json
+import pandas as pd
+# this will implement the search functionality
+
+
+def _project_aka_lookup(alias):
+    # this will lookup the project name based on the alias
+    conn = sqlite3.connect('backlogger.db')
+    c = conn.cursor()
+    c.execute(f"SELECT * FROM 'projects' WHERE alias = '{alias}'")
+    results = c.fetchall()
+    conn.close()
+    if len(results) > 1:
+        print("Error: multiple projects found with alias " + alias)
+    elif len(results) == 0:
+        raise Exception("Error: no project found with alias " + alias)
+    return results[0][0]
+
+
+def _db_lookup(db, ensemble, correlator_name, project=None, code=None, parameters=None, created_before=None, created_after=None, updated_before=None, updated_after=None, revision=None):
+    project_str = project
+
+    search_expr = f"SELECT * FROM 'backlogs' WHERE name = '{correlator_name}' AND ensemble = '{ensemble}'"
+    if project:
+        search_expr += f" AND project = '{project_str}'"
+    if code:
+        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}'"
+    print(search_expr)
+    conn = sqlite3.connect(db)
+    results = pd.read_sql(search_expr, conn)
+    conn.close()
+    return results
+
+
+def filter_results(results, **kwargs):
+    drops = []
+    for ind in range(len(results)):
+        result = results.iloc[ind]
+        if result['code'] == 'sfcf':
+            param = json.loads(result['parameters'])
+            if 'offset' in kwargs:
+                print("checking offset")
+                if kwargs.get('offset') != param['offset']:
+                    print("dropping")
+                    drops.append(ind)
+                    continue
+            if 'quark_masses' in kwargs:
+                quark_masses = kwargs['quark_masses']
+                if (quark_masses[0] != param['quarks'][0]['mass'] or quark_masses[1] != param['quarks'][1]['mass']) and (quark_masses[0] != param['quarks'][1]['mass'] or quark_masses[1] != param['quarks'][0]['mass']):
+                    drops.append(ind)
+                    continue
+            if 'quark_thetas' in kwargs:
+                quark_thetas = kwargs['quark_thetas']
+                if (quark_thetas[0] != param['quarks'][0]['thetas'] and quark_thetas[1] != param['quarks'][1]['thetas']) or (quark_thetas[0] != param['quarks'][1]['thetas'] and quark_thetas[1] != param['quarks'][0]['thetas']):
+                    drops.append(ind)
+                    continue
+            if 'wf1' in kwargs:
+                wf1 = kwargs['wf1']
+                if (wf1[0] != param['wf1'][0]) or (wf1[1][0] != param['wf1'][1][0]) or (wf1[1][1] != param['wf1'][1][1]):
+                    drops.append(ind)
+                    continue
+            if 'wf2' in kwargs:
+                wf2 = kwargs['wf2']
+                if (wf2[0] != param['wf2'][0]) or (wf2[1][0] != param['wf2'][1][0]) or (wf2[1][1] != param['wf2'][1][1]):
+                    drops.append(ind)
+                    continue
+
+    return results.drop(drops)
+
+
+def find_correlator(path, ensemble, correlator_name, project=None, code=None, parameters=None, created_before=None, created_after=None, updated_before=None, updated_after=None, revision=None, **kwargs):
+    db = path + '/backlogger.db'
+    if os.path.exists(db):
+        dl.get(db, dataset=path)
+    results = _db_lookup(db, ensemble, correlator_name, project, code=None, parameters=None, created_before=created_before, created_after=created_after, updated_before=updated_before, updated_after=updated_after, revision=revision)
+    results = filter_results(results, **kwargs)
+    print("Found " + str(len(results)) + " results")
+    return results