fix: upsert and insert with default_to_null boolean argument (#398)

Author: sapphire008

postgrest/_async/request_builder.py

@@ -303,6 +303,7 @@ def insert(
         count: Optional[CountMethod] = None,
         returning: ReturnMethod = ReturnMethod.representation,
         upsert: bool = False,
+        default_to_null: bool = True,
     ) -> AsyncQueryRequestBuilder[_ReturnT]:
         """Run an INSERT query.
 
@@ -311,6 +312,9 @@ def insert(
             count: The method to use to get the count of rows returned.
             returning: Either 'minimal' or 'representation'
             upsert: Whether the query should be an upsert.
+            default_to_null: Make missing fields default to `null`.
+                Otherwise, use the default value for the column.
+                Only applies for bulk inserts.
         Returns:
             :class:`AsyncQueryRequestBuilder`
         """
@@ -319,6 +323,7 @@ def insert(
             count=count,
             returning=returning,
             upsert=upsert,
+            default_to_null=default_to_null,
         )
         return AsyncQueryRequestBuilder[_ReturnT](
             self.session, self.path, method, headers, params, json
@@ -332,6 +337,7 @@ def upsert(
         returning: ReturnMethod = ReturnMethod.representation,
         ignore_duplicates: bool = False,
         on_conflict: str = "",
+        default_to_null: bool = True,
     ) -> AsyncQueryRequestBuilder[_ReturnT]:
         """Run an upsert (INSERT ... ON CONFLICT DO UPDATE) query.
 
@@ -341,6 +347,10 @@ def upsert(
             returning: Either 'minimal' or 'representation'
             ignore_duplicates: Whether duplicate rows should be ignored.
             on_conflict: Specified columns to be made to work with UNIQUE constraint.
+            default_to_null: Make missing fields default to `null`. Otherwise, use the
+                default value for the column. This only applies when inserting new rows,
+                not when merging with existing rows under `ignoreDuplicates: false`.
+                This also only applies when doing bulk upserts.
         Returns:
             :class:`AsyncQueryRequestBuilder`
         """
@@ -350,6 +360,7 @@ def upsert(
             returning=returning,
             ignore_duplicates=ignore_duplicates,
             on_conflict=on_conflict,
+            default_to_null=default_to_null,
         )
         return AsyncQueryRequestBuilder[_ReturnT](
             self.session, self.path, method, headers, params, json

postgrest/_sync/request_builder.py

@@ -34,7 +34,7 @@ def __init__(
         http_method: str,
         headers: Headers,
         params: QueryParams,
-        json: dict,
+        json: Union[dict, list],
     ) -> None:
         self.session = session
         self.path = path
@@ -289,7 +289,7 @@ def select(
             *columns: The names of the columns to fetch.
             count: The method to use to get the count of rows returned.
         Returns:
-            :class:`AsyncSelectRequestBuilder`
+            :class:`SyncSelectRequestBuilder`
         """
         method, params, headers, json = pre_select(*columns, count=count)
         return SyncSelectRequestBuilder[_ReturnT](
@@ -303,6 +303,7 @@ def insert(
         count: Optional[CountMethod] = None,
         returning: ReturnMethod = ReturnMethod.representation,
         upsert: bool = False,
+        default_to_null: bool = True,
     ) -> SyncQueryRequestBuilder[_ReturnT]:
         """Run an INSERT query.
 
@@ -311,14 +312,18 @@ def insert(
             count: The method to use to get the count of rows returned.
             returning: Either 'minimal' or 'representation'
             upsert: Whether the query should be an upsert.
+            default_to_null: Make missing fields default to `null`.
+                Otherwise, use the default value for the column.
+                Only applies for bulk inserts.
         Returns:
-            :class:`AsyncQueryRequestBuilder`
+            :class:`SyncQueryRequestBuilder`
         """
         method, params, headers, json = pre_insert(
             json,
             count=count,
             returning=returning,
             upsert=upsert,
+            default_to_null=default_to_null,
         )
         return SyncQueryRequestBuilder[_ReturnT](
             self.session, self.path, method, headers, params, json
@@ -332,6 +337,7 @@ def upsert(
         returning: ReturnMethod = ReturnMethod.representation,
         ignore_duplicates: bool = False,
         on_conflict: str = "",
+        default_to_null: bool = True,
     ) -> SyncQueryRequestBuilder[_ReturnT]:
         """Run an upsert (INSERT ... ON CONFLICT DO UPDATE) query.
 
@@ -341,15 +347,20 @@ def upsert(
             returning: Either 'minimal' or 'representation'
             ignore_duplicates: Whether duplicate rows should be ignored.
             on_conflict: Specified columns to be made to work with UNIQUE constraint.
+            default_to_null: Make missing fields default to `null`. Otherwise, use the
+                default value for the column. This only applies when inserting new rows,
+                not when merging with existing rows under `ignoreDuplicates: false`.
+                This also only applies when doing bulk upserts.
         Returns:
-            :class:`AsyncQueryRequestBuilder`
+            :class:`SyncQueryRequestBuilder`
         """
         method, params, headers, json = pre_upsert(
             json,
             count=count,
             returning=returning,
             ignore_duplicates=ignore_duplicates,
             on_conflict=on_conflict,
+            default_to_null=default_to_null,
         )
         return SyncQueryRequestBuilder[_ReturnT](
             self.session, self.path, method, headers, params, json
@@ -369,7 +380,7 @@ def update(
             count: The method to use to get the count of rows returned.
             returning: Either 'minimal' or 'representation'
         Returns:
-            :class:`AsyncFilterRequestBuilder`
+            :class:`SyncFilterRequestBuilder`
         """
         method, params, headers, json = pre_update(
             json,
@@ -392,7 +403,7 @@ def delete(
             count: The method to use to get the count of rows returned.
             returning: Either 'minimal' or 'representation'
         Returns:
-            :class:`AsyncFilterRequestBuilder`
+            :class:`SyncFilterRequestBuilder`
         """
         method, params, headers, json = pre_delete(
             count=count,

postgrest/base_request_builder.py

@@ -46,6 +46,12 @@ class QueryArgs(NamedTuple):
     json: Dict[Any, Any]
 
 
+def _unique_columns(json: List[Dict]):
+    unique_keys = {key for row in json for key in row.keys()}
+    columns = ",".join([f'"{k}"' for k in unique_keys])
+    return columns
+
+
 def pre_select(
     *columns: str,
     count: Optional[CountMethod] = None,
@@ -61,38 +67,51 @@ def pre_select(
 
 
 def pre_insert(
-    json: dict,
+    json: Union[dict, list],
     *,
     count: Optional[CountMethod],
     returning: ReturnMethod,
     upsert: bool,
+    default_to_null: bool = True,
 ) -> QueryArgs:
     prefer_headers = [f"return={returning}"]
     if count:
         prefer_headers.append(f"count={count}")
     if upsert:
         prefer_headers.append("resolution=merge-duplicates")
+    if not default_to_null:
+        prefer_headers.append("missing=default")
     headers = Headers({"Prefer": ",".join(prefer_headers)})
-    return QueryArgs(RequestMethod.POST, QueryParams(), headers, json)
+    # Adding 'columns' query parameters
+    query_params = {}
+    if isinstance(json, list):
+        query_params = {"columns": _unique_columns(json)}
+    return QueryArgs(RequestMethod.POST, QueryParams(query_params), headers, json)
 
 
 def pre_upsert(
-    json: dict,
+    json: Union[dict, list],
     *,
     count: Optional[CountMethod],
     returning: ReturnMethod,
     ignore_duplicates: bool,
     on_conflict: str = "",
+    default_to_null: bool = True,
 ) -> QueryArgs:
     query_params = {}
     prefer_headers = [f"return={returning}"]
     if count:
         prefer_headers.append(f"count={count}")
     resolution = "ignore" if ignore_duplicates else "merge"
     prefer_headers.append(f"resolution={resolution}-duplicates")
+    if not default_to_null:
+        prefer_headers.append("missing=default")
     headers = Headers({"Prefer": ",".join(prefer_headers)})
     if on_conflict:
         query_params["on_conflict"] = on_conflict
+    # Adding 'columns' query parameters
+    if isinstance(json, list):
+        query_params["columns"] = _unique_columns(json)
     return QueryArgs(RequestMethod.POST, QueryParams(query_params), headers, json)
 
 

tests/_async/test_request_builder.py

@@ -71,6 +71,31 @@ def test_insert_with_upsert(self, request_builder: AsyncRequestBuilder):
         assert builder.http_method == "POST"
         assert builder.json == {"key1": "val1"}
 
+    def test_upsert_with_default_single(self, request_builder: AsyncRequestBuilder):
+        builder = request_builder.upsert([{"key1": "val1"}], default_to_null=False)
+        assert builder.headers.get_list("prefer", True) == [
+            "return=representation",
+            "resolution=merge-duplicates",
+            "missing=default",
+        ]
+        assert builder.http_method == "POST"
+        assert builder.json == [{"key1": "val1"}]
+        assert builder.params.get("columns") == '"key1"'
+
+    def test_bulk_insert_using_default(self, request_builder: AsyncRequestBuilder):
+        builder = request_builder.insert(
+            [{"key1": "val1", "key2": "val2"}, {"key3": "val3"}], default_to_null=False
+        )
+        assert builder.headers.get_list("prefer", True) == [
+            "return=representation",
+            "missing=default",
+        ]
+        assert builder.http_method == "POST"
+        assert builder.json == [{"key1": "val1", "key2": "val2"}, {"key3": "val3"}]
+        assert set(builder.params["columns"].split(",")) == set(
+            '"key1","key2","key3"'.split(",")
+        )
+
     def test_upsert(self, request_builder: AsyncRequestBuilder):
         builder = request_builder.upsert({"key1": "val1"})
 
@@ -81,6 +106,21 @@ def test_upsert(self, request_builder: AsyncRequestBuilder):
         assert builder.http_method == "POST"
         assert builder.json == {"key1": "val1"}
 
+    def test_bulk_upsert_with_default(self, request_builder: AsyncRequestBuilder):
+        builder = request_builder.upsert(
+            [{"key1": "val1", "key2": "val2"}, {"key3": "val3"}], default_to_null=False
+        )
+        assert builder.headers.get_list("prefer", True) == [
+            "return=representation",
+            "resolution=merge-duplicates",
+            "missing=default",
+        ]
+        assert builder.http_method == "POST"
+        assert builder.json == [{"key1": "val1", "key2": "val2"}, {"key3": "val3"}]
+        assert set(builder.params["columns"].split(",")) == set(
+            '"key1","key2","key3"'.split(",")
+        )
+
 
 class TestUpdate:
     def test_update(self, request_builder: AsyncRequestBuilder):

tests/_sync/test_request_builder.py

@@ -71,6 +71,20 @@ def test_insert_with_upsert(self, request_builder: SyncRequestBuilder):
         assert builder.http_method == "POST"
         assert builder.json == {"key1": "val1"}
 
+    def test_bulk_insert_using_default(self, request_builder: SyncRequestBuilder):
+        builder = request_builder.insert(
+            [{"key1": "val1", "key2": "val2"}, {"key3": "val3"}], default_to_null=False
+        )
+        assert builder.headers.get_list("prefer", True) == [
+            "return=representation",
+            "missing=default",
+        ]
+        assert builder.http_method == "POST"
+        assert builder.json == [{"key1": "val1", "key2": "val2"}, {"key3": "val3"}]
+        assert set(builder.params["columns"].split(",")) == set(
+            '"key1","key2","key3"'.split(",")
+        )
+
     def test_upsert(self, request_builder: SyncRequestBuilder):
         builder = request_builder.upsert({"key1": "val1"})
 
@@ -81,6 +95,32 @@ def test_upsert(self, request_builder: SyncRequestBuilder):
         assert builder.http_method == "POST"
         assert builder.json == {"key1": "val1"}
 
+    def test_upsert_with_default_single(self, request_builder: SyncRequestBuilder):
+        builder = request_builder.upsert([{"key1": "val1"}], default_to_null=False)
+        assert builder.headers.get_list("prefer", True) == [
+            "return=representation",
+            "resolution=merge-duplicates",
+            "missing=default",
+        ]
+        assert builder.http_method == "POST"
+        assert builder.json == [{"key1": "val1"}]
+        assert builder.params.get("columns") == '"key1"'
+
+    def test_bulk_upsert_with_default(self, request_builder: SyncRequestBuilder):
+        builder = request_builder.upsert(
+            [{"key1": "val1", "key2": "val2"}, {"key3": "val3"}], default_to_null=False
+        )
+        assert builder.headers.get_list("prefer", True) == [
+            "return=representation",
+            "resolution=merge-duplicates",
+            "missing=default",
+        ]
+        assert builder.http_method == "POST"
+        assert builder.json == [{"key1": "val1", "key2": "val2"}, {"key3": "val3"}]
+        assert set(builder.params["columns"].split(",")) == set(
+            '"key1","key2","key3"'.split(",")
+        )
+
 
 class TestUpdate:
     def test_update(self, request_builder: SyncRequestBuilder):