From e5a1f568f4b0c2e2a0440528bf365d158979d1e0 Mon Sep 17 00:00:00 2001
From: "Huska, Matthew" <HuskaM@rki.de>
Date: Mon, 30 Oct 2023 13:55:55 +0100
Subject: [PATCH] If match returns no rows, still output column names

---
 lib/sonardb.py | 10 ++++++----
 sonar.py       | 19 ++++++++-----------
 2 files changed, 14 insertions(+), 15 deletions(-)

diff --git a/lib/sonardb.py b/lib/sonardb.py
index 167c4bb..4384307 100755
--- a/lib/sonardb.py
+++ b/lib/sonardb.py
@@ -1999,7 +1999,9 @@ def match(
             print("query: " + sql)
             print("vals: ", where_vals)
 
-        return self.cursor.execute(sql, where_vals).fetchall()
+        cursor = self.cursor.execute(sql, where_vals)
+        column_names = [description[0] for description in cursor.description]
+        return cursor.fetchall(), column_names
 
     # UPDATE DATA
 
@@ -3423,7 +3425,7 @@ def match(
             # print(include_lin)
 
             ########################
-            rows = dbm.match(
+            rows, column_names = dbm.match(
                 include_profiles,
                 exclude_profiles,
                 include_acc,
@@ -3479,7 +3481,7 @@ def match(
         elif count:
             return rows[0]["count"]
 
-        return rows
+        return rows, column_names
 
     # VALIDATION
 
@@ -3761,7 +3763,7 @@ def be_paranoid(self, acc, orig_seq, auto_delete=False, dbm=None):
                 )
 
             # frameshift checks
-            row = self.match(accessions=[acc], ambig=True, dbm=dbm)[0]
+            row = self.match(accessions=[acc], ambig=True, dbm=dbm)[0][0]
             fs = set()
             for dna_var in row["dna_profile"].split(" "):
                 if dna_var.strip() == "":
diff --git a/sonar.py b/sonar.py
index 2db042b..a1cd0ba 100755
--- a/sonar.py
+++ b/sonar.py
@@ -732,7 +732,7 @@ def match_genomes(
         frameshifts=0,
         tsv=False,
     ):
-        rows = self.db.match(
+        rows, column_names = self.db.match(
             include_profiles=include_profiles,
             exclude_profiles=exclude_profiles,
             accessions=accessions,
@@ -761,7 +761,7 @@ def match_genomes(
         if count:
             print(rows)
         else:
-            self.rows_to_csv(rows, na="*** no match ***", tsv=tsv)
+            self.rows_to_csv(rows, column_names, na="*** no match ***", tsv=tsv)
 
     def update_metadata(
         self,
@@ -917,17 +917,14 @@ def show_db_info(self):
                 spacer = " " * (maxlen - len(field))
                 print("   " + field + " information:" + spacer, f"{c} ({p:.{2}f}%)")
 
-    def rows_to_csv(self, rows, file=None, na="*** no data ***", tsv=False):
+    def rows_to_csv(self, rows, header, file=None, na="*** no data ***", tsv=False):
         if len(rows) == 0:
             print(na, file=sys.stderr)
-        else:
-            file = sys.stdout if file is None else open(file, "w")
-            sep = "\t" if tsv else ","
-            writer = csv.DictWriter(
-                file, rows[0].keys(), delimiter=sep, lineterminator=os.linesep
-            )
-            writer.writeheader()
-            writer.writerows(rows)
+        file = sys.stdout if file is None else open(file, "w")
+        sep = "\t" if tsv else ","
+        writer = csv.DictWriter(file, header, delimiter=sep, lineterminator=os.linesep)
+        writer.writeheader()
+        writer.writerows(rows)
 
     def get_db_size(self, decimal_places=3):
         size = os.path.getsize(self.dbfile)