feat(postgrest): implement max_affected method (#1222)

Co-authored-by: Claude <[email protected]>

Author: grdsdev

src/postgrest/src/postgrest/base_request_builder.py

@@ -526,6 +526,26 @@ def match(self: Self, query: Dict[str, Any]) -> Self:
 
         return updated_query
 
+    def max_affected(self: Self, value: int) -> Self:
+        """Set the maximum number of rows that can be affected by the query.
+
+        Only available in PostgREST v13+ and only works with PATCH and DELETE methods.
+
+        Args:
+            value: The maximum number of rows that can be affected
+        """
+        prefer_header = self.headers.get("Prefer", "")
+        if prefer_header:
+            if "handling=strict" not in prefer_header:
+                prefer_header += ",handling=strict"
+        else:
+            prefer_header = "handling=strict"
+
+        prefer_header += f",max-affected={value}"
+
+        self.headers["Prefer"] = prefer_header
+        return self
+
 
 class BaseSelectRequestBuilder(BaseFilterRequestBuilder[_ReturnT]):
     def __init__(

src/postgrest/tests/_async/test_filter_request_builder.py

@@ -241,3 +241,34 @@ def test_or_in_contain(filter_request_builder):
         str(builder.params)
         == "or=%28id.in.%285%2C6%2C7%29%2C+arraycol.cs.%7B%27a%27%2C%27b%27%7D%29"
     )
+
+
+def test_max_affected(filter_request_builder):
+    builder = filter_request_builder.max_affected(5)
+
+    assert builder.headers["prefer"] == "handling=strict,max-affected=5"
+
+
+def test_max_affected_with_existing_prefer_header(filter_request_builder):
+    # Set an existing prefer header
+    filter_request_builder.headers["prefer"] = "return=representation"
+    builder = filter_request_builder.max_affected(10)
+
+    assert (
+        builder.headers["prefer"]
+        == "return=representation,handling=strict,max-affected=10"
+    )
+
+
+def test_max_affected_with_existing_handling_strict(filter_request_builder):
+    # Set an existing prefer header with handling=strict
+    filter_request_builder.headers["prefer"] = "handling=strict,return=minimal"
+    builder = filter_request_builder.max_affected(3)
+
+    assert builder.headers["prefer"] == "handling=strict,return=minimal,max-affected=3"
+
+
+def test_max_affected_returns_self(filter_request_builder):
+    builder = filter_request_builder.max_affected(1)
+
+    assert builder is filter_request_builder

src/postgrest/tests/_async/test_request_builder.py

@@ -147,6 +147,15 @@ def test_update_with_count(self, request_builder: AsyncRequestBuilder):
         assert builder.http_method == "PATCH"
         assert builder.json == {"key1": "val1"}
 
+    def test_update_with_max_affected(self, request_builder: AsyncRequestBuilder):
+        builder = request_builder.update({"key1": "val1"}).max_affected(5)
+
+        assert "handling=strict" in builder.headers["prefer"]
+        assert "max-affected=5" in builder.headers["prefer"]
+        assert "return=representation" in builder.headers["prefer"]
+        assert builder.http_method == "PATCH"
+        assert builder.json == {"key1": "val1"}
+
 
 class TestDelete:
     def test_delete(self, request_builder: AsyncRequestBuilder):
@@ -166,6 +175,15 @@ def test_delete_with_count(self, request_builder: AsyncRequestBuilder):
         assert builder.http_method == "DELETE"
         assert builder.json == {}
 
+    def test_delete_with_max_affected(self, request_builder: AsyncRequestBuilder):
+        builder = request_builder.delete().max_affected(10)
+
+        assert "handling=strict" in builder.headers["prefer"]
+        assert "max-affected=10" in builder.headers["prefer"]
+        assert "return=representation" in builder.headers["prefer"]
+        assert builder.http_method == "DELETE"
+        assert builder.json == {}
+
 
 class TestTextSearch:
     def test_text_search(self, request_builder: AsyncRequestBuilder):

src/postgrest/tests/_sync/test_filter_request_builder.py

@@ -241,3 +241,34 @@ def test_or_in_contain(filter_request_builder):
         str(builder.params)
         == "or=%28id.in.%285%2C6%2C7%29%2C+arraycol.cs.%7B%27a%27%2C%27b%27%7D%29"
     )
+
+
+def test_max_affected(filter_request_builder):
+    builder = filter_request_builder.max_affected(5)
+
+    assert builder.headers["prefer"] == "handling=strict,max-affected=5"
+
+
+def test_max_affected_with_existing_prefer_header(filter_request_builder):
+    # Set an existing prefer header
+    filter_request_builder.headers["prefer"] = "return=representation"
+    builder = filter_request_builder.max_affected(10)
+
+    assert (
+        builder.headers["prefer"]
+        == "return=representation,handling=strict,max-affected=10"
+    )
+
+
+def test_max_affected_with_existing_handling_strict(filter_request_builder):
+    # Set an existing prefer header with handling=strict
+    filter_request_builder.headers["prefer"] = "handling=strict,return=minimal"
+    builder = filter_request_builder.max_affected(3)
+
+    assert builder.headers["prefer"] == "handling=strict,return=minimal,max-affected=3"
+
+
+def test_max_affected_returns_self(filter_request_builder):
+    builder = filter_request_builder.max_affected(1)
+
+    assert builder is filter_request_builder

src/postgrest/tests/_sync/test_request_builder.py

@@ -147,6 +147,15 @@ def test_update_with_count(self, request_builder: SyncRequestBuilder):
         assert builder.http_method == "PATCH"
         assert builder.json == {"key1": "val1"}
 
+    def test_update_with_max_affected(self, request_builder: SyncRequestBuilder):
+        builder = request_builder.update({"key1": "val1"}).max_affected(5)
+
+        assert "handling=strict" in builder.headers["prefer"]
+        assert "max-affected=5" in builder.headers["prefer"]
+        assert "return=representation" in builder.headers["prefer"]
+        assert builder.http_method == "PATCH"
+        assert builder.json == {"key1": "val1"}
+
 
 class TestDelete:
     def test_delete(self, request_builder: SyncRequestBuilder):
@@ -166,6 +175,15 @@ def test_delete_with_count(self, request_builder: SyncRequestBuilder):
         assert builder.http_method == "DELETE"
         assert builder.json == {}
 
+    def test_delete_with_max_affected(self, request_builder: SyncRequestBuilder):
+        builder = request_builder.delete().max_affected(10)
+
+        assert "handling=strict" in builder.headers["prefer"]
+        assert "max-affected=10" in builder.headers["prefer"]
+        assert "return=representation" in builder.headers["prefer"]
+        assert builder.http_method == "DELETE"
+        assert builder.json == {}
+
 
 class TestTextSearch:
     def test_text_search(self, request_builder: SyncRequestBuilder):