Add commit retries to insert table writes

src/tower/_tables.py

@@ -153,9 +153,21 @@ def rows_affected(self) -> RowsAffectedInformation:
         """
         return self._stats
 
-    def insert(self, data: pa.Table) -> TTable:
+    @staticmethod
+    def _validate_retry_args(max_retries: int, retry_delay_seconds: float) -> None:
+        if max_retries < 0:
+            raise ValueError("max_retries must be >= 0")
+        if retry_delay_seconds < 0:
+            raise ValueError("retry_delay_seconds must be >= 0")
+
+    def insert(
+        self,
+        data: pa.Table,
+        max_retries: int = 5,
+        retry_delay_seconds: float = 0.5,
+    ) -> TTable:
         """
-        Inserts new rows into the Iceberg table.
+        Inserts new rows into the Iceberg table. In case of commit conflicts, reloads the metadata and retries.
 
         This method appends the provided data to the table. The data must be provided as a
         PyArrow table with a schema that matches the table's schema. The operation is
@@ -164,10 +176,17 @@ def insert(self, data: pa.Table) -> TTable:
         Args:
             data (pa.Table): The data to insert into the table. The schema of this table
                 must match the schema of the target table.
+            max_retries (int): Maximum number of retry attempts on commit conflicts.
+                Defaults to 5.
+            retry_delay_seconds (float): Wait time in seconds between retries.
+                Defaults to 0.5 seconds.
 
         Returns:
             TTable: The table instance with the newly inserted rows, allowing for method chaining.
 
+        Raises:
+            CommitFailedException: If all retry attempts are exhausted.
+
         Example:
             >>> table = tables("my_table").load()
             >>> # Create a PyArrow table with new data
@@ -182,9 +201,25 @@ def insert(self, data: pa.Table) -> TTable:
             >>> stats = table.rows_affected()
             >>> print(f"Inserted {stats.inserts} rows")
         """
-        self._table.append(data)
-        self._stats.inserts += data.num_rows
-        return self
+        self._validate_retry_args(max_retries, retry_delay_seconds)
+
+        last_exception = None
+
+        for attempt in range(max_retries + 1):
+            try:
+                if attempt > 0:
+                    self._table.refresh()
+
+                self._table.append(data)
+                self._stats.inserts += data.num_rows
+                return self
+
+            except CommitFailedException as e:
+                last_exception = e
+                if attempt < max_retries:
+                    time.sleep(retry_delay_seconds)
+
+        raise last_exception
 
     def upsert(
         self,
@@ -239,6 +274,8 @@ def upsert(
             >>> print(f"Updated {stats.updates} rows")
             >>> print(f"Inserted {stats.inserts} rows")
         """
+        self._validate_retry_args(max_retries, retry_delay_seconds)
+
         last_exception = None
 
         for attempt in range(max_retries + 1):
@@ -268,9 +305,15 @@ def upsert(
 
         raise last_exception
 
-    def delete(self, filters: Union[str, List[pc.Expression]]) -> TTable:
+    def delete(
+        self,
+        filters: Union[str, List[pc.Expression]],
+        max_retries: int = 5,
+        retry_delay_seconds: float = 0.5,
+    ) -> TTable:
         """
         Deletes rows from the Iceberg table that match the specified filter conditions.
+        In case of commit conflicts, reloads the metadata and retries.
 
         This method removes rows from the table based on the provided filter expressions.
         The operation is always case-sensitive. Note that the number of deleted rows
@@ -282,10 +325,17 @@ def delete(self, filters: Union[str, List[pc.Expression]]) -> TTable:
                 - A single PyArrow compute expression
                 - A list of PyArrow compute expressions (combined with AND)
                 - A string expression
+            max_retries (int): Maximum number of retry attempts on commit conflicts.
+                Defaults to 5.
+            retry_delay_seconds (float): Wait time in seconds between retries.
+                Defaults to 0.5 seconds.
 
         Returns:
             TTable: The table instance with the deleted rows, allowing for method chaining.
 
+        Raises:
+            CommitFailedException: If all retry attempts are exhausted.
+
         Note:
             - The operation is always case-sensitive
             - The number of deleted rows cannot be tracked in the table statistics
@@ -303,22 +353,37 @@ def delete(self, filters: Union[str, List[pc.Expression]]) -> TTable:
             >>> # Delete rows using a string expression
             >>> table.delete("age > 30 AND department = 'IT'")
         """
+        self._validate_retry_args(max_retries, retry_delay_seconds)
 
         if isinstance(filters, list):
             # We need to convert the pc.Expression into PyIceberg
             next_filters = convert_pyarrow_expressions(filters)
             filters = next_filters
 
-        self._table.delete(
-            delete_filter=filters,
-            # We want this to always be the case. Not sure why you wouldn't?
-            case_sensitive=True,
-        )
+        last_exception = None
 
-        # NOTE: There is, unfortunately, no way to get the number of rows
-        # deleted besides comparing the two snapshots that were created.
+        for attempt in range(max_retries + 1):
+            try:
+                if attempt > 0:
+                    self._table.refresh()
+
+                self._table.delete(
+                    delete_filter=filters,
+                    # We want this to always be the case. Not sure why you wouldn't?
+                    case_sensitive=True,
+                )
+
+                # NOTE: There is, unfortunately, no way to get the number of rows
+                # deleted besides comparing the two snapshots that were created.
+
+                return self
 
-        return self
+            except CommitFailedException as e:
+                last_exception = e
+                if attempt < max_retries:
+                    time.sleep(retry_delay_seconds)
+
+        raise last_exception
 
     def schema(self) -> pa.Schema:
         """

tests/tower/test_tables.py

@@ -12,6 +12,7 @@
 import pyarrow as pa
 from pyiceberg.catalog.memory import InMemoryCatalog
 from pyiceberg.catalog.sql import SqlCatalog
+from pyiceberg.exceptions import CommitFailedException
 
 import concurrent.futures
 
@@ -321,6 +322,150 @@ def tracked_refresh():
     assert final_counter in [1, 2, 3, 4, 5]
 
 
+def test_insert_concurrent_writes_with_retry(sql_catalog):
+    """Test that concurrent inserts succeed with retry logic handling conflicts."""
+    schema = pa.schema(
+        [
+            pa.field("ticker", pa.string()),
+            pa.field("date", pa.string()),
+            pa.field("price", pa.float64()),
+        ]
+    )
+
+    ref = tower.tables("concurrent_insert_test", catalog=sql_catalog)
+    table = ref.create_if_not_exists(schema)
+
+    retry_count = {"value": 0}
+    retry_lock = threading.Lock()
+
+    def insert_ticker(ticker: str, price: float):
+        t = tower.tables("concurrent_insert_test", catalog=sql_catalog).load()
+
+        original_append = t._table.append
+        original_refresh = t._table.refresh
+        first_attempt = {"done": False}
+        refresh_called = {"value": False}
+
+        def failing_then_succeeding_append(data):
+            if not first_attempt["done"]:
+                first_attempt["done"] = True
+                raise CommitFailedException("Simulated concurrent write conflict")
+            return original_append(data)
+
+        def tracked_refresh():
+            refresh_called["value"] = True
+            with retry_lock:
+                retry_count["value"] += 1
+            return original_refresh()
+
+        t._table.append = failing_then_succeeding_append
+        t._table.refresh = tracked_refresh
+
+        data = pa.Table.from_pylist(
+            [{"ticker": ticker, "date": "2024-01-01", "price": price}],
+            schema=schema,
+        )
+        t.insert(data)
+        assert refresh_called["value"], "Expected refresh() to be called before retry"
+        return ticker
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
+        futures = [
+            executor.submit(insert_ticker, "AAPL", 150.0),
+            executor.submit(insert_ticker, "GOOGL", 250.0),
+            executor.submit(insert_ticker, "MSFT", 350.0),
+        ]
+        results = [f.result() for f in concurrent.futures.as_completed(futures)]
+
+    assert len(results) == 3
+    assert (
+        retry_count["value"] >= 3
+    ), "Expected at least one refresh per thread (one simulated failure each)"
+
+    final_table = tower.tables("concurrent_insert_test", catalog=sql_catalog).load()
+    df = final_table.read()
+
+    assert len(df) == 3
+
+    ticker_prices = {row["ticker"]: row["price"] for row in df.iter_rows(named=True)}
+
+    assert ticker_prices["AAPL"] == 150.0
+    assert ticker_prices["GOOGL"] == 250.0
+    assert ticker_prices["MSFT"] == 350.0
+
+
+def test_delete_concurrent_writes_with_retry(sql_catalog):
+    """Test that concurrent deletes succeed with retry logic handling conflicts."""
+    schema = pa.schema(
+        [
+            pa.field("ticker", pa.string()),
+            pa.field("date", pa.string()),
+            pa.field("price", pa.float64()),
+        ]
+    )
+
+    ref = tower.tables("concurrent_delete_test", catalog=sql_catalog)
+    table = ref.create_if_not_exists(schema)
+
+    initial_data = pa.Table.from_pylist(
+        [
+            {"ticker": "AAPL", "date": "2024-01-01", "price": 100.0},
+            {"ticker": "GOOGL", "date": "2024-01-01", "price": 200.0},
+            {"ticker": "MSFT", "date": "2024-01-01", "price": 300.0},
+        ],
+        schema=schema,
+    )
+    table.insert(initial_data)
+
+    retry_count = {"value": 0}
+    retry_lock = threading.Lock()
+
+    def delete_ticker(ticker: str):
+        t = tower.tables("concurrent_delete_test", catalog=sql_catalog).load()
+
+        original_delete = t._table.delete
+        original_refresh = t._table.refresh
+        first_attempt = {"done": False}
+        refresh_called = {"value": False}
+
+        def failing_then_succeeding_delete(**kwargs):
+            if not first_attempt["done"]:
+                first_attempt["done"] = True
+                raise CommitFailedException("Simulated concurrent write conflict")
+            return original_delete(**kwargs)
+
+        def tracked_refresh():
+            refresh_called["value"] = True
+            with retry_lock:
+                retry_count["value"] += 1
+            return original_refresh()
+
+        t._table.delete = failing_then_succeeding_delete
+        t._table.refresh = tracked_refresh
+
+        t.delete(filters=f"ticker = '{ticker}'")
+        assert refresh_called["value"], "Expected refresh() to be called before retry"
+        return ticker
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
+        futures = [
+            executor.submit(delete_ticker, "AAPL"),
+            executor.submit(delete_ticker, "GOOGL"),
+            executor.submit(delete_ticker, "MSFT"),
+        ]
+        results = [f.result() for f in concurrent.futures.as_completed(futures)]
+
+    assert len(results) == 3
+    assert (
+        retry_count["value"] >= 3
+    ), "Expected at least one refresh per thread (one simulated failure each)"
+
+    final_table = tower.tables("concurrent_delete_test", catalog=sql_catalog).load()
+    df = final_table.read()
+
+    assert len(df) == 0
+
+
 def test_delete_from_tables(in_memory_catalog):
     schema = pa.schema(
         [