Persist GridFS storage mode during waveform migration

python/mspasspy/db/database.py

@@ -3338,6 +3338,16 @@ def update_data(
         # Now handle update of sample data.  The gridfs method used here
         # handles that correctly based on the gridfs id.
         if mspass_object.live:
+            original_storage_mode = (
+                mspass_object["storage_mode"]
+                if "storage_mode" in mspass_object
+                else None
+            )
+            original_has_gridfs_id = "gridfs_id" in mspass_object
+            original_gridfs_id = (
+                mspass_object["gridfs_id"] if original_has_gridfs_id else None
+            )
+            file_to_gridfs_transition = original_storage_mode == "file"
             if "storage_mode" in mspass_object:
                 storage_mode = mspass_object["storage_mode"]
                 if not storage_mode == "gridfs":
@@ -3349,6 +3359,8 @@ def update_data(
                         ErrorSeverity.Complaint,
                     )
                     mspass_object["storage_mode"] = "gridfs"
+                    if file_to_gridfs_transition:
+                        update_record["storage_mode"] = "gridfs"
             else:
                 mspass_object.elog.log_error(
                     alg_name,
@@ -3359,7 +3371,10 @@ def update_data(
                 update_record["storage_mode"] = "gridfs"
             # Supplying the replacement id lets rollback remove a blob even
             # when GridFS writes it and then raises an exception.
-            staged_gridfs_id = ObjectId() if old_gridfs_id is not None else None
+            # A file-backed datum never treats a stale gridfs_id as its active
+            # sample reference; storage_mode remains authoritative.
+            active_old_gridfs_id = None if file_to_gridfs_transition else old_gridfs_id
+            staged_gridfs_id = ObjectId() if active_old_gridfs_id is not None else None
             new_gridfs_id = None
             elog_id = None
             old_elog_id = None
@@ -3395,7 +3410,6 @@ def update_data(
                             "to the waveform",
                             ErrorSeverity.Invalid,
                         )
-
                 if mspass_object.elog.size() > 0:
                     elog_id_name = self.database_schema.default_name("elog") + "_id"
                     # FIXME I think here we should check if elog_id field exists in the mspass_object
@@ -3420,13 +3434,15 @@ def update_data(
                     )
 
                 filter_ = {"_id": mspass_object["_id"]}
-                if old_gridfs_id is not None:
-                    filter_["gridfs_id"] = old_gridfs_id
+                if active_old_gridfs_id is not None:
+                    filter_["gridfs_id"] = active_old_gridfs_id
                 result = wf_collection.update_one(filter_, {"$set": update_record})
-                if old_gridfs_id is not None and result.matched_count != 1:
+                if (
+                    active_old_gridfs_id is not None or file_to_gridfs_transition
+                ) and result.matched_count != 1:
                     raise MsPASSError(
-                        "Database.update_data could not replace the expected "
-                        "GridFS reference",
+                        "Database.update_data could not commit the new GridFS "
+                        "reference",
                         ErrorSeverity.Invalid,
                     )
             except Exception as original_error:
@@ -3447,10 +3463,12 @@ def update_data(
                 except Exception as cleanup_error:
                     raise original_error from cleanup_error
                 finally:
-                    if old_gridfs_id is not None:
-                        mspass_object["gridfs_id"] = old_gridfs_id
+                    if original_has_gridfs_id:
+                        mspass_object["gridfs_id"] = original_gridfs_id
                     elif "gridfs_id" in mspass_object:
                         mspass_object.erase("gridfs_id")
+                    if file_to_gridfs_transition:
+                        mspass_object["storage_mode"] = original_storage_mode
                 raise
             # we may probably set the elog_id field in the mspass_object
             if elog_id:
@@ -3462,11 +3480,11 @@ def update_data(
                 self._reset_processing_history(
                     mspass_object, alg_name, alg_id, history_save_uuid
                 )
-            if old_gridfs_id is not None:
+            if active_old_gridfs_id is not None:
                 try:
                     gfsh = gridfs.GridFS(self)
-                    if gfsh.exists(old_gridfs_id):
-                        gfsh.delete(old_gridfs_id)
+                    if gfsh.exists(active_old_gridfs_id):
+                        gfsh.delete(active_old_gridfs_id)
                 except Exception as error:
                     mspass_object.elog.log_error(
                         alg_name,

python/tests/db/test_update_data_file_to_gridfs.py

@@ -0,0 +1,202 @@
+import os
+import uuid
+from types import SimpleNamespace
+from unittest.mock import patch
+
+import gridfs
+import numpy as np
+import pytest
+
+from mspasspy.ccore.seismic import (
+    DoubleVector,
+    Seismogram,
+    TimeReferenceType,
+    TimeSeries,
+)
+from mspasspy.ccore.utility import (
+    AtomicType,
+    ErrorSeverity,
+    MsPASSError,
+    dmatrix,
+)
+from mspasspy.db.client import DBClient
+from mspasspy.db.collection import Collection
+from mspasspy.db.database import Database
+
+
+@pytest.fixture
+def database():
+    uri = os.environ.get("MSPASS_TEST_MONGODB_URI", "mongodb://127.0.0.1:27017")
+    client = DBClient(uri, serverSelectionTimeoutMS=2000)
+    try:
+        client.admin.command("ping")
+    except Exception as error:
+        client.close()
+        pytest.skip(f"MongoDB is unavailable at {uri}: {error}")
+    name = "test_update_file_gridfs_" + uuid.uuid4().hex
+    database = Database(client, name)
+    try:
+        yield database
+    finally:
+        client.drop_database(name)
+        client.close()
+
+
+def make_timeseries(values):
+    datum = TimeSeries(len(values))
+    datum.data = DoubleVector(values)
+    datum.set_as_origin("test", "0", "0", AtomicType.TIMESERIES)
+    return finish_datum(datum, len(values))
+
+
+def make_seismogram(values):
+    datum = Seismogram(len(values))
+    samples = dmatrix(3, len(values))
+    for component in range(3):
+        for sample, value in enumerate(values):
+            samples[component, sample] = value + 10.0 * component
+    datum.data = samples
+    datum.set_as_origin("test", "0", "0", AtomicType.SEISMOGRAM)
+    return finish_datum(datum, len(values))
+
+
+def finish_datum(datum, npts):
+    datum.dt = 0.1
+    datum.t0 = 10.0
+    datum.tref = TimeReferenceType.UTC
+    datum.set_live()
+    datum["npts"] = npts
+    datum["sampling_rate"] = 10.0
+    datum["delta"] = 0.1
+    datum["calib"] = 1.0
+    return datum
+
+
+def sample_array(datum):
+    return np.asarray(datum.data)
+
+
+def storage_reference(document):
+    keys = ("storage_mode", "dir", "dfile", "foff", "format", "nbytes", "gridfs_id")
+    return {key: document[key] for key in keys if key in document}
+
+
+def add_stale_gridfs_reference(database, collection, datum):
+    stale_gridfs_id = gridfs.GridFS(database).put(b"unreachable sample data")
+    database[collection].update_one(
+        {"_id": datum["_id"]}, {"$set": {"gridfs_id": stale_gridfs_id}}
+    )
+    datum["gridfs_id"] = stale_gridfs_id
+    return stale_gridfs_id
+
+
+@pytest.mark.parametrize(
+    "factory,collection",
+    [
+        (make_timeseries, "wf_TimeSeries"),
+        (make_seismogram, "wf_Seismogram"),
+    ],
+)
+@pytest.mark.parametrize("stale_gridfs_reference", [False, True])
+def test_file_to_gridfs_update_round_trips_new_samples(
+    database, tmp_path, factory, collection, stale_gridfs_reference
+):
+    original = factory([1.0, 2.0, 3.0, 4.0])
+    datum = database.save_data(
+        original,
+        mode="promiscuous",
+        storage_mode="file",
+        dir=str(tmp_path),
+        dfile="samples.bin",
+        save_history=False,
+        return_data=True,
+    )
+    stale_gridfs_id = None
+    if stale_gridfs_reference:
+        stale_gridfs_id = add_stale_gridfs_reference(database, collection, datum)
+    replacement = factory([5.0, 6.0, 7.0, 8.0])
+    datum.data = replacement.data
+
+    result = database.update_data(datum, mode="promiscuous")
+
+    document = database[collection].find_one({"_id": result["_id"]})
+    assert document["storage_mode"] == "gridfs"
+    assert document["gridfs_id"] == result["gridfs_id"]
+    if stale_gridfs_reference:
+        assert document["gridfs_id"] != stale_gridfs_id
+        assert gridfs.GridFS(database).exists(stale_gridfs_id)
+    assert gridfs.GridFS(database).exists(document["gridfs_id"])
+    reread = database.read_data(result["_id"], collection=collection)
+    assert reread.live
+    np.testing.assert_array_equal(sample_array(reread), sample_array(replacement))
+
+
+@pytest.mark.parametrize(
+    "factory,collection",
+    [
+        (make_timeseries, "wf_TimeSeries"),
+        (make_seismogram, "wf_Seismogram"),
+    ],
+)
+@pytest.mark.parametrize("failure_mode", ["exception", "unmatched"])
+@pytest.mark.parametrize("stale_gridfs_reference", [False, True])
+def test_failed_waveform_update_keeps_file_readable_and_removes_new_gridfs_data(
+    database,
+    tmp_path,
+    factory,
+    collection,
+    failure_mode,
+    stale_gridfs_reference,
+):
+    original = factory([1.0, 2.0, 3.0, 4.0])
+    original_samples = sample_array(original).copy()
+    datum = database.save_data(
+        original,
+        mode="promiscuous",
+        storage_mode="file",
+        dir=str(tmp_path),
+        dfile="samples.bin",
+        save_history=False,
+        return_data=True,
+    )
+    stale_gridfs_id = None
+    if stale_gridfs_reference:
+        stale_gridfs_id = add_stale_gridfs_reference(database, collection, datum)
+    replacement = factory([5.0, 6.0, 7.0, 8.0])
+    datum.data = replacement.data
+    original_document = database[collection].find_one({"_id": datum["_id"]})
+    original_reference = storage_reference(original_document)
+    original_gridfs_files = list(database["fs.files"].find({}))
+    original_gridfs_chunks = list(database["fs.chunks"].find({}))
+    original_update_one = Collection.update_one
+    injected_error = RuntimeError("injected waveform update failure")
+
+    def fail_waveform_reference_update(self, query, update, *args, **kwargs):
+        if "gridfs_id" in update.get("$set", {}):
+            if failure_mode == "exception":
+                raise injected_error
+            return SimpleNamespace(matched_count=0)
+        return original_update_one(self, query, update, *args, **kwargs)
+
+    with patch.object(Collection, "update_one", fail_waveform_reference_update):
+        if failure_mode == "exception":
+            with pytest.raises(RuntimeError) as error:
+                database.update_data(datum, mode="promiscuous")
+            assert error.value is injected_error
+        else:
+            with pytest.raises(MsPASSError) as error:
+                database.update_data(datum, mode="promiscuous")
+            assert error.value.severity == ErrorSeverity.Invalid
+
+    assert datum["storage_mode"] == "file"
+    if stale_gridfs_reference:
+        assert datum["gridfs_id"] == stale_gridfs_id
+    else:
+        assert "gridfs_id" not in datum
+    document = database[collection].find_one({"_id": datum["_id"]})
+    assert storage_reference(document) == original_reference
+    assert list(database["fs.files"].find({})) == original_gridfs_files
+    assert list(database["fs.chunks"].find({})) == original_gridfs_chunks
+    reread = database.read_data(datum["_id"], collection=collection)
+    assert reread.live
+    np.testing.assert_array_equal(sample_array(reread), original_samples)