Repository: cogent3/cogent3
Pull request #1443 (2023-05-17 07:03:42)

MAINT: improve data store state consistency


Patch:
diff --git a/src/cogent3/app/composable.py b/src/cogent3/app/composable.py
index dc29d5e98..40bc53dfe 100644
--- a/src/cogent3/app/composable.py
+++ b/src/cogent3/app/composable.py
@@ -754,7 +754,9 @@ def _apply_to(
         input_id = id_from_source(input_id)
         if input_id in inputs:
             raise ValueError("non-unique identifier detected in data")
-        if input_id in self.data_store:  # todo write a test
+        if input_id in self.data_store:
+            # we are assuming that this query returns True only when
+            # an input_id is completed, we will not hit this if not_completed
             continue
         inputs[input_id] = m
 
diff --git a/src/cogent3/app/data_store_new.py b/src/cogent3/app/data_store_new.py
index d98a25a4e..d95ca8049 100644
--- a/src/cogent3/app/data_store_new.py
+++ b/src/cogent3/app/data_store_new.py
@@ -12,6 +12,7 @@ from abc import ABC, abstractmethod
 from collections import defaultdict
 from enum import Enum
 from functools import singledispatch
+from io import TextIOWrapper
 from os import PathLike
 from pathlib import Path
 from typing import Iterator, Optional, Union
@@ -250,7 +251,7 @@ class DataStoreABC(ABC):
         )
 
     @abstractmethod
-    def drop_not_completed(self):
+    def drop_not_completed(self, *, unique_id: Optional[str] = None) -> None:
         ...
 
     def validate(self) -> TabularType:
@@ -428,16 +429,25 @@ class DataStoreDirectory(DataStoreABC):
             data = infile.read()
         return data
 
-    def drop_not_completed(self):
+    def drop_not_completed(self, *, unique_id: str = "") -> None:
+        unique_id = unique_id.replace(f".{self.suffix}", "")
+        unique_id = f"{unique_id}.json" if unique_id else unique_id
         nc_dir = self.source / _NOT_COMPLETED_TABLE
         md5_dir = self.source / _MD5_TABLE
-        for file in nc_dir.glob("*.json"):
+        for m in list(self.not_completed):
+            if unique_id and not m.unique_id.endswith(unique_id):
+                continue
+
+            file = nc_dir / Path(m.unique_id).name
             file.unlink()
             md5_file = md5_dir / f"{file.stem}.txt"
             md5_file.unlink()
-        Path(self.source / _NOT_COMPLETED_TABLE).rmdir()
-        # reset _not_completed list to force not_completed function to make it again
-        self._not_completed = []
+            self.not_completed.remove(m)
+
+        if not unique_id:
+            Path(self.source / _NOT_COMPLETED_TABLE).rmdir()
+            # reset _not_completed list to force not_completed function to make it again
+            self._not_completed = []
 
     @property
     def logs(self) -> list[DataMember]:
@@ -494,7 +504,7 @@ class DataStoreDirectory(DataStoreABC):
         if suffix != "log" and unique_id in self:
             return None
 
-        with open_(self.source / subdir / unique_id, mode="w") as out:
+        with open_(self.source / subdir / unique_id, mode="w", newline="\n") as out:
             out.write(data)
 
         if subdir == _LOG_TABLE:
@@ -515,14 +525,45 @@ class DataStoreDirectory(DataStoreABC):
         return member
 
     def write(self, *, unique_id: str, data: str) -> DataMember:
+        """writes a completed record ending with .suffix
+
+        Parameters
+        ----------
+        unique_id
+            unique identifier
+        data
+            text data to be written
+
+        Returns
+        -------
+        a member for this record
+
+        Notes
+        -----
+        Drops any not-completed member corresponding to this identifier
+        """
         member = self._write(
             subdir="", unique_id=unique_id, suffix=self.suffix, data=data
         )
+        self.drop_not_completed(unique_id=unique_id)
         if member is not None:
             self._completed.append(member)
         return member
 
     def write_not_completed(self, *, unique_id: str, data: str) -> DataMember:
+        """writes a not completed record as json
+
+        Parameters
+        ----------
+        unique_id
+            unique identifier
+        data
+            text data to be written
+
+        Returns
+        -------
+        a member for this record
+        """
         (self.source / _NOT_COMPLETED_TABLE).mkdir(parents=True, exist_ok=True)
         member = self._write(
             subdir=_NOT_COMPLETED_TABLE, unique_id=unique_id, suffix="json", data=data
@@ -593,7 +634,8 @@ class ReadOnlyDataStoreZipped(DataStoreABC):
         unique_id = str(pathlib.Path(self.source.stem, unique_id)).replace("\\", "/")
         with zipfile.ZipFile(self.source) as archive:
             record = archive.open(unique_id)
-            return record.read().decode("utf8")
+            record = TextIOWrapper(record, encoding="latin-1")
+            return record.read()
 
     def _iter_matches(self, subdir: str, pattern: str) -> Iterator[PathLike]:
         with zipfile.ZipFile(self._source) as archive:
@@ -665,7 +707,7 @@ class ReadOnlyDataStoreZipped(DataStoreABC):
             m = DataMember(data_store=self, unique_id=str(md5_dir / name.name))
             return m.read()
 
-    def drop_not_completed(self):
+    def drop_not_completed(self, *, unique_id: Optional[str] = None) -> None:
         raise TypeError("zip data stores are read only")
 
     def write(self, *, unique_id: str, data: StrOrBytes) -> None:
diff --git a/src/cogent3/app/sqlite_data_store.py b/src/cogent3/app/sqlite_data_store.py
index 879114c3e..c5ff99556 100644
--- a/src/cogent3/app/sqlite_data_store.py
+++ b/src/cogent3/app/sqlite_data_store.py
@@ -20,9 +20,11 @@ from cogent3.app.data_store_new import (
     DataMember,
     DataMemberABC,
     DataStoreABC,
+    DataStoreDirectory,
     Mode,
     StrOrBytes,
 )
+from cogent3.util.misc import extend_docstring_from
 
 
 _RESULT_TABLE = "results"
@@ -265,8 +267,14 @@ class DataStoreSqlite(DataStoreABC):
 
         return DataMember(data_store=self, unique_id=unique_id)
 
-    def drop_not_completed(self) -> None:
-        self.db.execute(f"DELETE FROM {_RESULT_TABLE} WHERE is_completed=0")
+    def drop_not_completed(self, *, unique_id: str = "") -> None:
+        if not unique_id:
+            cmnd = f"DELETE FROM {_RESULT_TABLE} WHERE is_completed=?"
+            vals = (0,)
+        else:
+            cmnd = f"DELETE FROM {_RESULT_TABLE} WHERE is_completed=? AND record_id=?"
+            vals = (0, unique_id)
+        self.db.execute(cmnd, vals)
         self._not_completed = []
 
     @property
@@ -322,11 +330,15 @@ class DataStoreSqlite(DataStoreABC):
 
         return
 
+    @extend_docstring_from(DataStoreDirectory.write)
     def write(self, *, unique_id: str, data: StrOrBytes) -> DataMemberABC:
         if unique_id.startswith(_RESULT_TABLE):
             unique_id = Path(unique_id).name
 
         super().write(unique_id=unique_id, data=data)
+
+        self.drop_not_completed(unique_id=unique_id)
+
         member = self._write(
             table_name=_RESULT_TABLE,
             unique_id=unique_id,
@@ -339,6 +351,7 @@ class DataStoreSqlite(DataStoreABC):
             self._completed.append(member)
         return member
 
+    @extend_docstring_from(DataStoreDirectory.write_log)
     def write_log(self, *, unique_id: str, data: StrOrBytes) -> None:
         if unique_id.startswith(_LOG_TABLE):
             unique_id = Path(unique_id).name
@@ -348,6 +361,7 @@ class DataStoreSqlite(DataStoreABC):
             table_name=_LOG_TABLE, unique_id=unique_id, data=data, is_completed=False
         )
 
+    @extend_docstring_from(DataStoreDirectory.write_not_completed)
     def write_not_completed(self, *, unique_id: str, data: StrOrBytes) -> DataMemberABC:
         if unique_id.startswith(_RESULT_TABLE):
             unique_id = Path(unique_id).name


Test patch:
diff --git a/tests/test_app/test_data_store_new.py b/tests/test_app/test_data_store_new.py
index 05cee6c7c..c6ad9892a 100644
--- a/tests/test_app/test_data_store_new.py
+++ b/tests/test_app/test_data_store_new.py
@@ -481,6 +481,16 @@ def test_write_if_member_exists(full_dstore, write_dir):  # new changes
     assert got == expect
 
 
+def test_write_success_replaces_not_completed(full_dstore):
+    """correctly write content"""
+    nc = full_dstore.not_completed[0].unique_id
+    data = full_dstore.completed[0].read()
+    new_id = Path(nc.replace(".json", f".{full_dstore.suffix}")).name
+    num = len(full_dstore)
+    full_dstore.write(unique_id=new_id, data=data)
+    assert len(full_dstore) == num
+
+
 @pytest.mark.parametrize("klass", (str, Path))
 def test_get_data_source_attr(klass):
     """handles case where input has source attribute string object or pathlib object"""
@@ -595,7 +605,7 @@ def test_zipped_ro_ioerror():
 
 
 def _get_member_data(members) -> dict:
-    return {m.unique_id: m.read() for m in members}
+    return {m.unique_id: m.read().splitlines() for m in members}
 
 
 def test_zipped_ro_basic(zipped_basic, ro_dstore):
@@ -647,8 +657,6 @@ def test_zipped_md5(zipped_full, full_dstore):
 def test_zipped_contains(dstore, request):
     full_dstore = request.getfixturevalue(dstore)
     assert "formattest.fasta" in full_dstore
-    # following line does not work, should it?
-    # assert "id_2.json" in full_dstore
 
 
 @pytest.mark.parametrize("dstore", ("full_dstore", "zipped_full"))
diff --git a/tests/test_app/test_sqlite_data_store.py b/tests/test_app/test_sqlite_data_store.py
index a97c65e1c..228d33e5c 100644
--- a/tests/test_app/test_sqlite_data_store.py
+++ b/tests/test_app/test_sqlite_data_store.py
@@ -285,6 +285,15 @@ def test_read(full_dstore_sqlite):
         assert isinstance(r.read(), str)
 
 
+def test_write_success_replaces_not_completed(full_dstore_sqlite):
+    """correctly write content"""
+    new_id = full_dstore_sqlite.not_completed[0].unique_id
+    data = full_dstore_sqlite.completed[0].read()
+    num = len(full_dstore_sqlite)
+    full_dstore_sqlite.write(unique_id=new_id, data=data)
+    assert len(full_dstore_sqlite) == num
+
+
 def test_read_log(sql_dstore, log_data):
     """correctly read content"""
     sql_dstore.write_log(unique_id="brca1.fasta", data=log_data)