update inclusive behaviour and test

Author: chinmay-bhat

pyiceberg/table/__init__.py

@@ -1303,10 +1303,15 @@ def snapshot_by_name(self, name: str) -> Optional[Snapshot]:
             return self.snapshot_by_id(ref.snapshot_id)
         return None
 
-    def snapshot_at_or_before_timestamp(self, timestamp_ms: int) -> Optional[Snapshot]:
-        """Get the snapshot that was current at or right before the given timestamp, or None if there is no matching snapshot."""
+    def snapshot_as_of_timestamp(self, timestamp_ms: int, inclusive: bool = True) -> Optional[Snapshot]:
+        """Get the snapshot that was current as of or right before the given timestamp, or None if there is no matching snapshot.
+
+        Args:
+            timestamp_ms: Find snapshot that was current at/before this timestamp
+            inclusive: Includes timestamp_ms in search when True. Excludes timestamp_ms when False
+        """
         for log_entry in reversed(self.history()):
-            if log_entry.timestamp_ms <= timestamp_ms:
+            if (inclusive and log_entry.timestamp_ms <= timestamp_ms) or log_entry.timestamp_ms < timestamp_ms:
                 return self.snapshot_by_id(log_entry.snapshot_id)
         return None
 

tests/table/test_init.py

@@ -205,15 +205,16 @@ def test_snapshot_by_id(table_v2: Table) -> None:
 
 
 def test_snapshot_by_timestamp(table_v2: Table) -> None:
-    assert table_v2.snapshot_at_or_before_timestamp(1555100955771) == Snapshot(
-        snapshot_id=3055729675574597004,
-        parent_snapshot_id=3051729675574597004,
-        sequence_number=1,
-        timestamp_ms=1555100955770,
-        manifest_list="s3://a/b/2.avro",
-        summary=Summary(operation=Operation.APPEND),
-        schema_id=1,
+    assert table_v2.snapshot_as_of_timestamp(1515100955770) == Snapshot(
+        snapshot_id=3051729675574597004,
+        parent_snapshot_id=None,
+        sequence_number=0,
+        timestamp_ms=1515100955770,
+        manifest_list='s3://a/b/1.avro',
+        summary=Summary(Operation.APPEND),
+        schema_id=None,
     )
+    assert table_v2.snapshot_as_of_timestamp(1515100955770, inclusive=False) is None
 
 
 def test_current_ancestors(table_v2: Table) -> None: