update_annotation gets db_file to make the change

mehrdadbakhtiari · Jun 19, 2019 · 798e211 · 798e211
1 parent 4123989
commit 798e211
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/advntr/models.py b/advntr/models.py
@@ -176,8 +176,10 @@ def update_trained_score_in_database(vntr_id, scaled_recruitment_score):
     db.close()
 
 
-def update_gene_name_and_annotation_in_database(vntr_id, gene_name, annotation):
-    db = sqlite3.connect(settings.TRAINED_MODELS_DB)
+def update_gene_name_and_annotation_in_database(vntr_id, gene_name, annotation, db_file=None):
+    if db_file is None:
+        db_file = settings.TRAINED_MODELS_DB
+    db = sqlite3.connect(db_file)
     cursor = db.cursor()
     cursor.execute('''UPDATE vntrs SET gene_name=?, annotation=? WHERE id=?''', (gene_name, annotation, vntr_id))
     db.commit()