From d574f7e3a987882a89ec47716b513b826dc96af6 Mon Sep 17 00:00:00 2001 From: Utkrsh Date: Thu, 27 Aug 2026 00:28:07 +0530 Subject: [PATCH] fix: allow single response for insert select --- .../src/postgrest/_async/request_builder.py | 26 +++++++++---------- .../src/postgrest/_sync/request_builder.py | 26 +++++++++---------- .../tests/_async/test_request_builder.py | 19 +++++++++++++- .../tests/_sync/test_request_builder.py | 19 +++++++++++++- 4 files changed, 62 insertions(+), 28 deletions(-) diff --git a/src/postgrest/src/postgrest/_async/request_builder.py b/src/postgrest/src/postgrest/_async/request_builder.py index f5c34e5f..cfa8d144 100644 --- a/src/postgrest/src/postgrest/_async/request_builder.py +++ b/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: diff --git a/src/postgrest/src/postgrest/_sync/request_builder.py b/src/postgrest/src/postgrest/_sync/request_builder.py index df502b69..bf3ef709 100644 --- a/src/postgrest/src/postgrest/_sync/request_builder.py +++ b/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: diff --git a/src/postgrest/tests/_async/test_request_builder.py b/src/postgrest/tests/_async/test_request_builder.py index 9cdc0baa..76829055 100644 --- a/src/postgrest/tests/_async/test_request_builder.py +++ b/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 diff --git a/src/postgrest/tests/_sync/test_request_builder.py b/src/postgrest/tests/_sync/test_request_builder.py index 435f8ab5..efcff1f3 100644 --- a/src/postgrest/tests/_sync/test_request_builder.py +++ b/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