Skip to content

Commit 4f6e9d9

Browse files
feat: add csv() modifier (#316)
* fix: cast to correct type * feat: add csv() modifier * chore: export SingleRequestBuilder * chore: write tests for csv() * 'Refactored by Sourcery' (#317) Co-authored-by: Sourcery AI <> --------- Co-authored-by: sourcery-ai[bot] <58596630+sourcery-ai[bot]@users.noreply.github.com>
1 parent 5ae0f99 commit 4f6e9d9

File tree

6 files changed

+92
-5
lines changed

6 files changed

+92
-5
lines changed

postgrest/__init__.py

+2
Original file line numberDiff line numberDiff line change
@@ -10,13 +10,15 @@
1010
AsyncQueryRequestBuilder,
1111
AsyncRequestBuilder,
1212
AsyncSelectRequestBuilder,
13+
AsyncSingleRequestBuilder,
1314
)
1415
from ._sync.client import SyncPostgrestClient
1516
from ._sync.request_builder import (
1617
SyncFilterRequestBuilder,
1718
SyncQueryRequestBuilder,
1819
SyncRequestBuilder,
1920
SyncSelectRequestBuilder,
21+
SyncSingleRequestBuilder,
2022
)
2123
from .base_request_builder import APIResponse
2224
from .constants import DEFAULT_POSTGREST_CLIENT_HEADERS

postgrest/_async/request_builder.py

+12
Original file line numberDiff line numberDiff line change
@@ -252,6 +252,18 @@ def text_search(
252252
session=self.session, # type: ignore
253253
)
254254

255+
def csv(self) -> AsyncSingleRequestBuilder[str]:
256+
"""Specify that the query must retrieve data as a single CSV string."""
257+
self.headers["Accept"] = "text/csv"
258+
return AsyncSingleRequestBuilder[str](
259+
session=self.session, # type: ignore
260+
path=self.path,
261+
http_method=self.http_method,
262+
headers=self.headers,
263+
params=self.params,
264+
json=self.json,
265+
)
266+
255267

256268
class AsyncRequestBuilder(Generic[_ReturnT]):
257269
def __init__(self, session: AsyncClient, path: str) -> None:

postgrest/_sync/request_builder.py

+12
Original file line numberDiff line numberDiff line change
@@ -252,6 +252,18 @@ def text_search(
252252
session=self.session, # type: ignore
253253
)
254254

255+
def csv(self) -> SyncSingleRequestBuilder[str]:
256+
"""Specify that the query must retrieve data as a single CSV string."""
257+
self.headers["Accept"] = "text/csv"
258+
return SyncSingleRequestBuilder[str](
259+
session=self.session, # type: ignore
260+
path=self.path,
261+
http_method=self.http_method,
262+
headers=self.headers,
263+
params=self.params,
264+
json=self.json,
265+
)
266+
255267

256268
class SyncRequestBuilder(Generic[_ReturnT]):
257269
def __init__(self, session: SyncClient, path: str) -> None:

postgrest/base_request_builder.py

+6-3
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@
3434
from pydantic import validator as field_validator
3535

3636
from .types import CountMethod, Filters, RequestMethod, ReturnMethod
37-
from .utils import AsyncClient, SyncClient, sanitize_param
37+
from .utils import AsyncClient, SyncClient, get_origin_and_cast, sanitize_param
3838

3939

4040
class QueryArgs(NamedTuple):
@@ -200,8 +200,11 @@ class SingleAPIResponse(APIResponse[_ReturnT], Generic[_ReturnT]):
200200
def from_http_request_response(
201201
cls: Type[Self], request_response: RequestResponse
202202
) -> Self:
203-
data = request_response.json()
204203
count = cls._get_count_from_http_request_response(request_response)
204+
try:
205+
data = request_response.json()
206+
except JSONDecodeError:
207+
data = request_response.text if len(request_response.text) > 0 else []
205208
return cls[_ReturnT](data=data, count=count) # type: ignore
206209

207210
@classmethod
@@ -420,7 +423,7 @@ def __init__(
420423
# Generic[T] is an instance of typing._GenericAlias, so doing Generic[T].__init__
421424
# tries to call _GenericAlias.__init__ - which is the wrong method
422425
# The __origin__ attribute of the _GenericAlias is the actual class
423-
BaseFilterRequestBuilder[_ReturnT].__origin__.__init__(
426+
get_origin_and_cast(BaseFilterRequestBuilder[_ReturnT]).__init__(
424427
self, session, headers, params
425428
)
426429

tests/_async/test_request_builder.py

+30-1
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import pytest
44
from httpx import Request, Response
55

6-
from postgrest import AsyncRequestBuilder
6+
from postgrest import AsyncRequestBuilder, AsyncSingleRequestBuilder
77
from postgrest.base_request_builder import APIResponse, SingleAPIResponse
88
from postgrest.types import CountMethod
99
from postgrest.utils import AsyncClient
@@ -36,6 +36,12 @@ def test_select_with_count(self, request_builder: AsyncRequestBuilder):
3636
assert builder.http_method == "HEAD"
3737
assert builder.json == {}
3838

39+
def test_select_as_csv(self, request_builder: AsyncRequestBuilder):
40+
builder = request_builder.select("*").csv()
41+
42+
assert builder.headers["Accept"] == "text/csv"
43+
assert isinstance(builder, AsyncSingleRequestBuilder)
44+
3945

4046
class TestInsert:
4147
def test_insert(self, request_builder: AsyncRequestBuilder):
@@ -146,6 +152,11 @@ def test_explain_options(self, request_builder: AsyncRequestBuilder):
146152
)
147153

148154

155+
@pytest.fixture
156+
def csv_api_response() -> str:
157+
return "id,name\n1,foo\n"
158+
159+
149160
@pytest.fixture
150161
def api_response_with_error() -> Dict[str, Any]:
151162
return {
@@ -281,6 +292,15 @@ def request_response_with_single_data(
281292
)
282293

283294

295+
@pytest.fixture
296+
def request_response_with_csv_data(csv_api_response: str) -> Response:
297+
return Response(
298+
status_code=200,
299+
text=csv_api_response,
300+
request=Request(method="GET", url="http://example.com"),
301+
)
302+
303+
284304
class TestApiResponse:
285305
def test_response_raises_when_api_error(
286306
self, api_response_with_error: Dict[str, Any]
@@ -374,3 +394,12 @@ def test_single_from_http_request_response_constructor(
374394
assert isinstance(result.data, dict)
375395
assert result.data == single_api_response
376396
assert result.count == 2
397+
398+
def test_single_with_csv_data(
399+
self, request_response_with_csv_data: Response, csv_api_response: str
400+
):
401+
result = SingleAPIResponse.from_http_request_response(
402+
request_response_with_csv_data
403+
)
404+
assert isinstance(result.data, str)
405+
assert result.data == csv_api_response

tests/_sync/test_request_builder.py

+30-1
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
import pytest
44
from httpx import Request, Response
55

6-
from postgrest import SyncRequestBuilder
6+
from postgrest import SyncRequestBuilder, SyncSingleRequestBuilder
77
from postgrest.base_request_builder import APIResponse, SingleAPIResponse
88
from postgrest.types import CountMethod
99
from postgrest.utils import SyncClient
@@ -36,6 +36,12 @@ def test_select_with_count(self, request_builder: SyncRequestBuilder):
3636
assert builder.http_method == "HEAD"
3737
assert builder.json == {}
3838

39+
def test_select_as_csv(self, request_builder: SyncRequestBuilder):
40+
builder = request_builder.select("*").csv()
41+
42+
assert builder.headers["Accept"] == "text/csv"
43+
assert isinstance(builder, SyncSingleRequestBuilder)
44+
3945

4046
class TestInsert:
4147
def test_insert(self, request_builder: SyncRequestBuilder):
@@ -146,6 +152,11 @@ def test_explain_options(self, request_builder: SyncRequestBuilder):
146152
)
147153

148154

155+
@pytest.fixture
156+
def csv_api_response() -> str:
157+
return "id,name\n1,foo\n"
158+
159+
149160
@pytest.fixture
150161
def api_response_with_error() -> Dict[str, Any]:
151162
return {
@@ -281,6 +292,15 @@ def request_response_with_single_data(
281292
)
282293

283294

295+
@pytest.fixture
296+
def request_response_with_csv_data(csv_api_response: str) -> Response:
297+
return Response(
298+
status_code=200,
299+
text=csv_api_response,
300+
request=Request(method="GET", url="http://example.com"),
301+
)
302+
303+
284304
class TestApiResponse:
285305
def test_response_raises_when_api_error(
286306
self, api_response_with_error: Dict[str, Any]
@@ -374,3 +394,12 @@ def test_single_from_http_request_response_constructor(
374394
assert isinstance(result.data, dict)
375395
assert result.data == single_api_response
376396
assert result.count == 2
397+
398+
def test_single_with_csv_data(
399+
self, request_response_with_csv_data: Response, csv_api_response: str
400+
):
401+
result = SingleAPIResponse.from_http_request_response(
402+
request_response_with_csv_data
403+
)
404+
assert isinstance(result.data, str)
405+
assert result.data == csv_api_response

0 commit comments

Comments
 (0)