Impl and tests

Author: ognyanstoimenov

python/arcticdb/version_store/library.py

@@ -1281,17 +1281,17 @@ def delete_staged_data(self, symbol: str):
     def finalize_staged_data(
         self,
         symbol: str,
-        mode: Optional[StagedDataFinalizeMethod] = StagedDataFinalizeMethod.WRITE,
+        mode: Optional[Union[StagedDataFinalizeMethod, str]] = StagedDataFinalizeMethod.WRITE,
         prune_previous_versions: bool = False,
         metadata: Any = None,
         validate_index = True,
         delete_staged_data_on_failure: bool = False
     ) -> VersionedItem:
         """
         Finalizes staged data, making it available for reads. All staged segments must be ordered and non-overlapping.
-        ``finalize_staged_data`` is less time consuming than ``sort_and_finalize_staged_data``.
+        ``finalize_staged_data`` is less time-consuming than ``sort_and_finalize_staged_data``.
 
-        If ``mode`` is ``StagedDataFinalizeMethod.APPEND`` the index of the first row of the new segment must be equal to or greater
+        If ``mode`` is ``StagedDataFinalizeMethod.APPEND`` or ``APPEND`` the index of the first row of the new segment must be equal to or greater
         than the index of the last row in the existing data.
 
         If ``Static Schema`` is used all staged block must have matching schema (same column names, same dtype, same column ordering)
@@ -1310,9 +1310,9 @@ def finalize_staged_data(
         symbol : `str`
             Symbol to finalize data for.
 
-        mode : `StagedDataFinalizeMethod`, default=StagedDataFinalizeMethod.WRITE
-            Finalize mode. Valid options are WRITE or APPEND. Write collects the staged data and writes them to a
-            new version. Append collects the staged data and appends them to the latest version.
+        mode : Union[`StagedDataFinalizeMethod`, str], default=StagedDataFinalizeMethod.WRITE
+            Finalize mode. Valid options are StagedDataFinalizeMethod.WRITE or StagedDataFinalizeMethod.APPEND. Write collects the staged data and writes them to a
+            new version. Append collects the staged data and appends them to the latest version. Also accepts "WRITE" and "APPEND".
         prune_previous_versions: bool, default=False
             Removes previous (non-snapshotted) versions from the database.
         metadata : Any, default=None
@@ -1383,9 +1383,12 @@ def finalize_staged_data(
         2024-01-03    3
         2024-01-04    4
         """
+        if not (mode is None or isinstance(mode, StagedDataFinalizeMethod) or mode == "WRITE" or mode == "APPEND"):
+            raise ArcticInvalidApiUsageException("mode must be a StagedDataFinalizeMethod enum or 'WRITE'/'APPEND'")
+
         return self._nvs.compact_incomplete(
             symbol,
-            append=mode == StagedDataFinalizeMethod.APPEND,
+            append=mode == StagedDataFinalizeMethod.APPEND or mode == "APPEND",
             convert_int_to_float=False,
             metadata=metadata,
             prune_previous_version=prune_previous_versions,

python/tests/unit/arcticdb/version_store/test_api.py

@@ -8,10 +8,12 @@
 import time
 
 from pandas import Timestamp
+from unittest.mock import MagicMock
 import pytest
 
 from arcticdb.exceptions import NoSuchVersionException, NoDataFoundException
 from arcticdb.util.test import distinct_timestamps
+from arcticdb.version_store.library import StagedDataFinalizeMethod, ArcticInvalidApiUsageException
 
 
 def test_read_descriptor(lmdb_version_store, one_col_df):
@@ -89,3 +91,38 @@ def test_get_num_rows_pickled(lmdb_version_store):
     symbol = "test_get_num_rows_pickled"
     lmdb_version_store.write(symbol, 1)
     assert lmdb_version_store.get_num_rows(symbol) is None
+
+
+@pytest.mark.parametrize(
+    "input_mode, expected_append",
+    [
+        ("WRITE", False),
+        ("APPEND", True),
+        (None, False),
+        (StagedDataFinalizeMethod.APPEND, True),
+        (StagedDataFinalizeMethod.WRITE, False),
+        ([], None),
+        {"write", None},
+    ],
+)
+def test_finalize_staged_data(arctic_library_lmdb, input_mode, expected_append):
+    symbol = "sym"
+    arctic_library_lmdb._nvs = MagicMock()
+
+    if expected_append is None:
+        with pytest.raises(ArcticInvalidApiUsageException):
+            arctic_library_lmdb.finalize_staged_data(symbol, input_mode)
+        return
+
+    arctic_library_lmdb.finalize_staged_data(symbol, input_mode)
+
+    default_args = {
+        "convert_int_to_float": False,
+        "metadata": None,
+        "prune_previous_version": False,
+        "validate_index": True,
+        "delete_staged_data_on_failure": False,
+    }
+
+    arctic_library_lmdb._nvs.compact_incomplete.assert_called_once_with(symbol, append=expected_append, **default_args)
+