Allow single response builder after insert select

src/postgrest/src/postgrest/_async/request_builder.py

@@ -71,6 +71,19 @@ def select(self: QueryBuilderT, *columns: str) -> QueryBuilderT:
             self.request.headers["Prefer"] = "return=representation"
         return self
 
+    def single(self) -> AsyncSingleRequestBuilder:
+        """Specify that the query will only return a single row in response.
+
+        .. caution::
+            The API will raise an error if the query returned more than one row.
+        """
+        self.request.headers["Accept"] = "application/vnd.pgrst.object+json"
+        return AsyncSingleRequestBuilder(self.request)
+
+    def maybe_single(self) -> AsyncMaybeSingleRequestBuilder:
+        """Retrieves at most one row from the result. Result must be at most one row (e.g. using `eq` on a UNIQUE column), otherwise this will result in an error."""
+        return AsyncMaybeSingleRequestBuilder(self.request)
+
     def retry(self, enabled: bool) -> Self:
         self.request.retry_enabled = enabled
         return self
@@ -209,19 +222,6 @@ def __init__(self, request: ReqConfig) -> None:
         BaseSelectRequestBuilder.__init__(self, request)
         AsyncQueryRequestBuilder.__init__(self, request)
 
-    def single(self) -> AsyncSingleRequestBuilder:
-        """Specify that the query will only return a single row in response.
-
-        .. caution::
-            The API will raise an error if the query returned more than one row.
-        """
-        self.request.headers["Accept"] = "application/vnd.pgrst.object+json"
-        return AsyncSingleRequestBuilder(self.request)
-
-    def maybe_single(self) -> AsyncMaybeSingleRequestBuilder:
-        """Retrieves at most one row from the result. Result must be at most one row (e.g. using `eq` on a UNIQUE column), otherwise this will result in an error."""
-        return AsyncMaybeSingleRequestBuilder(self.request)
-
     def text_search(
         self, column: str, query: str, options: dict[str, Any] = {}
     ) -> AsyncQueryRequestBuilder:

src/postgrest/src/postgrest/_sync/request_builder.py

@@ -71,6 +71,19 @@ def select(self: QueryBuilderT, *columns: str) -> QueryBuilderT:
             self.request.headers["Prefer"] = "return=representation"
         return self
 
+    def single(self) -> SyncSingleRequestBuilder:
+        """Specify that the query will only return a single row in response.
+
+        .. caution::
+            The API will raise an error if the query returned more than one row.
+        """
+        self.request.headers["Accept"] = "application/vnd.pgrst.object+json"
+        return SyncSingleRequestBuilder(self.request)
+
+    def maybe_single(self) -> SyncMaybeSingleRequestBuilder:
+        """Retrieves at most one row from the result. Result must be at most one row (e.g. using `eq` on a UNIQUE column), otherwise this will result in an error."""
+        return SyncMaybeSingleRequestBuilder(self.request)
+
     def retry(self, enabled: bool) -> Self:
         self.request.retry_enabled = enabled
         return self
@@ -209,19 +222,6 @@ def __init__(self, request: ReqConfig) -> None:
         BaseSelectRequestBuilder.__init__(self, request)
         SyncQueryRequestBuilder.__init__(self, request)
 
-    def single(self) -> SyncSingleRequestBuilder:
-        """Specify that the query will only return a single row in response.
-
-        .. caution::
-            The API will raise an error if the query returned more than one row.
-        """
-        self.request.headers["Accept"] = "application/vnd.pgrst.object+json"
-        return SyncSingleRequestBuilder(self.request)
-
-    def maybe_single(self) -> SyncMaybeSingleRequestBuilder:
-        """Retrieves at most one row from the result. Result must be at most one row (e.g. using `eq` on a UNIQUE column), otherwise this will result in an error."""
-        return SyncMaybeSingleRequestBuilder(self.request)
-
     def text_search(
         self, column: str, query: str, options: dict[str, Any] = {}
     ) -> SyncQueryRequestBuilder:

src/postgrest/tests/_async/test_request_builder.py

@@ -4,7 +4,11 @@
 from httpx import AsyncClient, Headers, QueryParams, Request, Response
 from yarl import URL
 
-from postgrest import AsyncRequestBuilder, AsyncSingleRequestBuilder
+from postgrest import (
+    AsyncMaybeSingleRequestBuilder,
+    AsyncRequestBuilder,
+    AsyncSingleRequestBuilder,
+)
 from postgrest._async.request_builder import RequestConfig
 from postgrest.base_request_builder import APIResponse, SingleAPIResponse
 from postgrest.types import JSON, CountMethod, ReturnMethod
@@ -140,6 +144,19 @@ def test_insert_with_select_forces_representation(
             "return=representation",
         ]
 
+    def test_insert_with_select_single(self, request_builder: AsyncRequestBuilder):
+        builder = request_builder.insert({"key1": "val1"}).select("id").single()
+
+        assert builder.request.params["select"] == "id"
+        assert builder.request.headers["Accept"] == "application/vnd.pgrst.object+json"
+        assert isinstance(builder, AsyncSingleRequestBuilder)
+
+    def test_insert_with_select_maybe_single(self, request_builder: AsyncRequestBuilder):
+        builder = request_builder.insert({"key1": "val1"}).select("id").maybe_single()
+
+        assert builder.request.params["select"] == "id"
+        assert isinstance(builder, AsyncMaybeSingleRequestBuilder)
+
     def test_bulk_upsert_with_default(self, request_builder: AsyncRequestBuilder):
         builder = request_builder.upsert(
             [{"key1": "val1", "key2": "val2"}, {"key3": "val3"}], default_to_null=False

src/postgrest/tests/_sync/test_request_builder.py

@@ -4,7 +4,11 @@
 from httpx import Client, Headers, QueryParams, Request, Response
 from yarl import URL
 
-from postgrest import SyncRequestBuilder, SyncSingleRequestBuilder
+from postgrest import (
+    SyncMaybeSingleRequestBuilder,
+    SyncRequestBuilder,
+    SyncSingleRequestBuilder,
+)
 from postgrest._async.request_builder import RequestConfig
 from postgrest.base_request_builder import APIResponse, SingleAPIResponse
 from postgrest.types import JSON, CountMethod, ReturnMethod
@@ -140,6 +144,19 @@ def test_insert_with_select_forces_representation(
             "return=representation",
         ]
 
+    def test_insert_with_select_single(self, request_builder: SyncRequestBuilder):
+        builder = request_builder.insert({"key1": "val1"}).select("id").single()
+
+        assert builder.request.params["select"] == "id"
+        assert builder.request.headers["Accept"] == "application/vnd.pgrst.object+json"
+        assert isinstance(builder, SyncSingleRequestBuilder)
+
+    def test_insert_with_select_maybe_single(self, request_builder: SyncRequestBuilder):
+        builder = request_builder.insert({"key1": "val1"}).select("id").maybe_single()
+
+        assert builder.request.params["select"] == "id"
+        assert isinstance(builder, SyncMaybeSingleRequestBuilder)
+
     def test_bulk_upsert_with_default(self, request_builder: SyncRequestBuilder):
         builder = request_builder.upsert(
             [{"key1": "val1", "key2": "val2"}, {"key3": "val3"}], default_to_null=False