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