From 215aafbd2dcc13460462ec4faa9ca920e8e7814f Mon Sep 17 00:00:00 2001
From: Eric Neilsen <neilsen@fnal.gov>
Date: Wed, 13 Nov 2024 14:05:39 -0800
Subject: [PATCH] improve support for using data from consdb

---
 schedview/collect/consdb.py |  4 +++-
 schedview/collect/visits.py |  1 -
 schedview/compute/maf.py    | 23 ++++++++++++++++++++---
 3 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/schedview/collect/consdb.py b/schedview/collect/consdb.py
index 4354d75b..5302a141 100644
--- a/schedview/collect/consdb.py
+++ b/schedview/collect/consdb.py
@@ -16,7 +16,9 @@ def read_consdb(
     # right dtypes, so use ObservationArray()[0:0] instead to make the
     # empty recarray with the correct dtypes.
     visit_records: np.recarray = (
-        consdb_visits.opsim.to_records() if len(consdb_visits.opsim) > 0 else ObservationArray()[0:0]
+        consdb_visits.merged_opsim_consdb.to_records()
+        if len(consdb_visits.opsim) > 0
+        else ObservationArray()[0:0]
     )
     for stacker in stackers:
         visit_records = stacker.run(visit_records)
diff --git a/schedview/collect/visits.py b/schedview/collect/visits.py
index 41d6cff4..e9a425b3 100644
--- a/schedview/collect/visits.py
+++ b/schedview/collect/visits.py
@@ -18,7 +18,6 @@
 NIGHT_STACKERS = [
     maf.HourAngleStacker(),
     maf.stackers.ObservationStartDatetime64Stacker(),
-    maf.stackers.TeffStacker(),
     maf.stackers.OverheadStacker(),
     maf.stackers.DayObsISOStacker(),
 ]
diff --git a/schedview/compute/maf.py b/schedview/compute/maf.py
index 8c360fcd..636e7be4 100644
--- a/schedview/compute/maf.py
+++ b/schedview/compute/maf.py
@@ -1,3 +1,4 @@
+import sqlite3
 from pathlib import Path
 from tempfile import TemporaryDirectory
 
@@ -10,11 +11,27 @@
 
 def _visits_to_opsim(visits, opsim):
     # Take advantage of the schema migration code in SchemaConverter to make
-    # sure we have the correct column names
-
+    # sure we have up to date column names.
     schema_converter = SchemaConverter()
     obs = schema_converter.opsimdf2obs(visits)
-    schema_converter.obs2opsim(obs, filename=opsim)
+    updated_opsim = schema_converter.obs2opsim(obs)
+
+    # We can't just use the update opsim as is, because it might drop columns
+    # we want. Instead, merge the results back into the visits passed to us.
+    restored_columns = set(visits.columns) - set(updated_opsim.columns)
+    restored_columns.add("observationId")
+    merged_opsim = updated_opsim.merge(
+        visits.loc[:, list(restored_columns)], on="observationId", suffixes=("", "_orig")
+    )
+
+    # If the round trip change actually changes values in an existing column
+    # without changing the names, the merge might not work correctly. If
+    # this happens, the default "inner join" performed by DataFrame.merge
+    # will drop visits. Double check that this hasn't happened.
+    assert len(merged_opsim) == len(visits)
+
+    with sqlite3.connect(opsim) as con:
+        merged_opsim.to_sql("observations", con)
 
 
 def compute_metric(visits, metric_bundle):